Kulexiu/KulexiuForStudent/KulexiuForStudent/Common/MediaMerge/AudioMerge/MergeAudioAnimation/KSSpectrumView.m

//
//  KSSpectrumView.m
//  KulexiuSchoolStudent
//
//  Created by 王智 on 2024/7/29.
//

#import "KSSpectrumView.h"
#import <Accelerate/Accelerate.h>

@interface KSSpectrumView ()

@property (nonatomic, strong) CAGradientLayer *combinedGradientLayer;

@property (nonatomic, strong) CAShapeLayer *combinedMaskLayer;

@property (nonatomic, assign) CGFloat xIncrement;

@end

@implementation KSSpectrumView

- (instancetype)initWithFrame:(CGRect)frame {
    self = [super initWithFrame:frame];
    if (self) {
        [self setupView];
    }
    return self;
}


- (instancetype)initWithCoder:(NSCoder *)aDecoder {
    self = [super initWithCoder:aDecoder];
    if (self) {
        [self setupView];
    }
    return self;
}

- (void)setupView {
    self.barWidth = 3.0;
    self.space = 1.0;
    self.xIncrement = self.barWidth + self.space;
    
    self.combinedGradientLayer = [CAGradientLayer layer];
    self.combinedGradientLayer.colors = @[
        (__bridge id)HexRGBAlpha(0xffffff, 1.0f).CGColor,
        (__bridge id)HexRGBAlpha(0xffffff, 1.0f).CGColor
    ];
    self.combinedGradientLayer.locations = @[@0.6, @1.0];
    [self.layer addSublayer:self.combinedGradientLayer];
    
    self.combinedMaskLayer = [CAShapeLayer layer];
}

- (void)resetLayer {
    self.spectra = [NSArray array];
}

- (void)setSpectra:(NSArray<NSArray<NSNumber *> *> *)spectra {
    _spectra = spectra;
    if (spectra) {
        CGFloat viewHeight = self.bounds.size.height;
        CGFloat viewWidth = self.bounds.size.width;
        @weakObj(self);
        dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
            @strongObj(self);
            self.isModify = YES;
            NSUInteger spectraCount = [spectra[0] count];
            NSMutableArray<NSNumber *> *combinedSpectrum = [NSMutableArray arrayWithCapacity:spectraCount];
            // 取两个声道数据中的最大值
            for (NSUInteger i = 0; i < spectraCount; i++) {
                NSNumber *leftAmplitude = spectra[0][i];
                NSNumber *rightAmplitude = spectra.count > 1 ? spectra[1][i] : @0;
                CGFloat maxAmplitude = MAX(leftAmplitude.floatValue, rightAmplitude.floatValue);
                [combinedSpectrum addObject:@(maxAmplitude)];
            }
            
            CGFloat middleY = viewHeight / 2.0;
            CGFloat barHeight = (viewHeight) / 2.0;
            CGFloat cornerRadius = viewWidth / 2.0f;
            UIBezierPath *combinedPath = [UIBezierPath bezierPath];
            
            // Left channel
            for (NSUInteger i = 0; i < spectraCount; i++) {
                CGFloat x = i * self.xIncrement + self.space;
                CGFloat amplitudeValue = combinedSpectrum[i].floatValue;
                CGFloat height = amplitudeValue * barHeight;
                CGFloat y = middleY - height/2.0; // Centered vertically
                
                CGRect rect = CGRectMake(x, y, self.barWidth, height);
                UIBezierPath *barPath = [UIBezierPath bezierPathWithRoundedRect:rect cornerRadius:cornerRadius];
                [combinedPath appendPath:barPath];
            }
            
            dispatch_async(dispatch_get_main_queue(), ^{
                self.combinedMaskLayer.path = combinedPath.CGPath;
                self.combinedGradientLayer.frame = CGRectMake(0, 0, viewWidth, viewHeight);
                self.combinedGradientLayer.mask = self.combinedMaskLayer;
                self.isModify = NO;
            });
        });
    }
}

- (CGFloat)translateAmplitudeToYPosition:(float)amplitude {
    CGFloat barHeight = amplitude * self.bounds.size.height;
    return self.bounds.size.height - barHeight;
}

- (void)dealloc {
    NSLog(@"---- KSSpectrumView dealloc");
}

- (NSArray<NSNumber *> *)computeFFTWithPCMBuffer:(float *)pcmBuffer frameCount:(NSUInteger)frameCount {
    NSUInteger log2n = log2(frameCount); // FFT 长度必须是 2 的幂
    NSUInteger fftSize = 1 << log2n;    // 计算实际 FFT 点数
    
    // 创建 FFT 设置
    FFTSetup fftSetup = vDSP_create_fftsetup(log2n, FFT_RADIX2);
    if (!fftSetup) {
        NSLog(@"FFT Setup failed");
        return nil;
    }

    // 输入数据必须是复数形式，创建 SplitComplex 存储数据
    DSPSplitComplex splitComplex;
    float *real = malloc(sizeof(float) * fftSize / 2);
    float *imag = malloc(sizeof(float) * fftSize / 2);
    splitComplex.realp = real;
    splitComplex.imagp = imag;
    
    // 填充输入数据并进行窗口处理（如 Hanning 窗口）
    vDSP_ctoz((DSPComplex *)pcmBuffer, 2, &splitComplex, 1, fftSize / 2);
    vDSP_hann_window(real, fftSize / 2, vDSP_HANN_NORM);

    // 执行 FFT
    vDSP_fft_zrip(fftSetup, &splitComplex, 1, log2n, FFT_FORWARD);

    // 计算幅值（模长）
    float *amplitudes = malloc(sizeof(float) * fftSize / 2);
    vDSP_zvmags(&splitComplex, 1, amplitudes, 1, fftSize / 2);
    
    // 转换为 dB（可选）
    float scale = 1.0 / (2.0 * fftSize); // 归一化
    vDSP_vsmul(amplitudes, 1, &scale, amplitudes, 1, fftSize / 2);
    vDSP_vdbcon(amplitudes, 1, &scale, amplitudes, 1, fftSize / 2, 0);

    // 转换结果为 NSArray
    NSMutableArray<NSNumber *> *result = [NSMutableArray arrayWithCapacity:fftSize / 2];
    for (NSUInteger i = 0; i < fftSize / 2; i++) {
        [result addObject:@(amplitudes[i])];
    }

    // 清理内存
    free(real);
    free(imag);
    free(amplitudes);
    vDSP_destroy_fftsetup(fftSetup);

    return result;
}
/*
// Only override drawRect: if you perform custom drawing.
// An empty implementation adversely affects performance during animation.
- (void)drawRect:(CGRect)rect {
    // Drawing code
}
*/

@end