//
//  Downmix.m
//  Cog
//
//  Created by Christopher Snowhill on 2/05/22.
//  Copyright 2022 __LoSnoCo__. All rights reserved.
//

#import "Downmix.h"

#import "Logging.h"

static const float STEREO_DOWNMIX[8-2][8][2]={
    /*3.0*/
    {
        {0.5858F,0.0F},{0.0F,0.5858F},{0.4142F,0.4142F}
    },
    /*quadrophonic*/
    {
        {0.4226F,0.0F},{0.0F,0.4226F},{0.366F,0.2114F},{0.2114F,0.336F}
    },
    /*5.0*/
    {
        {0.651F,0.0F},{0.0F,0.651F},{0.46F,0.46F},{0.5636F,0.3254F},
        {0.3254F,0.5636F}
    },
    /*5.1*/
    {
        {0.529F,0.0F},{0.0F,0.529F},{0.3741F,0.3741F},{0.3741F,0.3741F},{0.4582F,0.2645F},
        {0.2645F,0.4582F}
    },
    /*6.1*/
    {
        {0.4553F,0.0F},{0.0F,0.4553F},{0.322F,0.322F},{0.322F,0.322F},{0.2788F,0.2788F},
        {0.3943F,0.2277F},{0.2277F,0.3943F}
    },
    /*7.1*/
    {
        {0.3886F,0.0F},{0.0F,0.3886F},{0.2748F,0.2748F},{0.2748F,0.2748F},{0.3366F,0.1943F},
        {0.1943F,0.3366F},{0.3366F,0.1943F},{0.1943F,0.3366F}
    }
};

static void downmix_to_stereo(const float * inBuffer, int channels, float * outBuffer, size_t count)
{
    if (channels >= 3 && channels <= 8)
    for (size_t i = 0; i < count; ++i)
    {
        float left = 0, right = 0;
        for (int j = 0; j < channels; ++j)
        {
            left += inBuffer[i * channels + j] * STEREO_DOWNMIX[channels - 3][j][0];
            right += inBuffer[i * channels + j] * STEREO_DOWNMIX[channels - 3][j][1];
        }
        outBuffer[i * 2 + 0] = left;
        outBuffer[i * 2 + 1] = right;
    }
}

static void downmix_to_mono(const float * inBuffer, int channels, float * outBuffer, size_t count)
{
    float tempBuffer[count * 2];
    if (channels >= 3 && channels <= 8)
    {
        downmix_to_stereo(inBuffer, channels, tempBuffer, count);
        inBuffer = tempBuffer;
        channels = 2;
    }
    float invchannels = 1.0 / (float)channels;
    for (size_t i = 0; i < count; ++i)
    {
        float sample = 0;
        for (int j = 0; j < channels; ++j)
        {
            sample += inBuffer[i * channels + j];
        }
        outBuffer[i] = sample * invchannels;
    }
}

static void upmix(const float * inBuffer, int inchannels, float * outBuffer, int outchannels, size_t count)
{
    for (ssize_t i = 0; i < count; ++i)
    {
        if (inchannels == 1 && outchannels == 2)
        {
            // upmix mono to stereo
            float sample = inBuffer[i];
            outBuffer[i * 2 + 0] = sample;
            outBuffer[i * 2 + 1] = sample;
        }
        else if (inchannels == 1 && outchannels == 4)
        {
            // upmix mono to quad
            float sample = inBuffer[i];
            outBuffer[i * 4 + 0] = sample;
            outBuffer[i * 4 + 1] = sample;
            outBuffer[i * 4 + 2] = 0;
            outBuffer[i * 4 + 3] = 0;
        }
        else if (inchannels == 1 && (outchannels == 3 || outchannels >= 5))
        {
            // upmix mono to center channel
            float sample = inBuffer[i];
            outBuffer[i * outchannels + 2] = sample;
            for (int j = 0; j < 2; ++j)
            {
                outBuffer[i * outchannels + j] = 0;
            }
            for (int j = 3; j < outchannels; ++j)
            {
                outBuffer[i * outchannels + j] = 0;
            }
        }
        else if (inchannels == 4 && outchannels >= 5)
        {
            float fl = inBuffer[i * 4 + 0];
            float fr = inBuffer[i * 4 + 1];
            float bl = inBuffer[i * 4 + 2];
            float br = inBuffer[i * 4 + 3];
            const int skipclfe = (outchannels == 5) ? 1 : 2;
            outBuffer[i * outchannels + 0] = fl;
            outBuffer[i * outchannels + 1] = fr;
            outBuffer[i * outchannels + skipclfe + 2] = bl;
            outBuffer[i * outchannels + skipclfe + 3] = br;
            for (int j = 0; j < skipclfe; ++j)
            {
                outBuffer[i * outchannels + 2 + j] = 0;
            }
            for (int j = 4 + skipclfe; j < outchannels; ++j)
            {
                outBuffer[i * outchannels + j] = 0;
            }
        }
        else if (inchannels == 5 && outchannels >= 6)
        {
            float fl = inBuffer[i * 5 + 0];
            float fr = inBuffer[i * 5 + 1];
            float c = inBuffer[i * 5 + 2];
            float bl = inBuffer[i * 5 + 3];
            float br = inBuffer[i * 5 + 4];
            outBuffer[i * outchannels + 0] = fl;
            outBuffer[i * outchannels + 1] = fr;
            outBuffer[i * outchannels + 2] = c;
            outBuffer[i * outchannels + 3] = 0;
            outBuffer[i * outchannels + 4] = bl;
            outBuffer[i * outchannels + 5] = br;
            for (int j = 6; j < outchannels; ++j)
            {
                outBuffer[i * outchannels + j] = 0;
            }
        }
        else if (inchannels == 7 && outchannels == 8)
        {
            float fl = inBuffer[i * 7 + 0];
            float fr = inBuffer[i * 7 + 1];
            float c = inBuffer[i * 7 + 2];
            float lfe = inBuffer[i * 7 + 3];
            float sl = inBuffer[i * 7 + 4];
            float sr = inBuffer[i * 7 + 5];
            float bc = inBuffer[i * 7 + 6];
            outBuffer[i * 8 + 0] = fl;
            outBuffer[i * 8 + 1] = fr;
            outBuffer[i * 8 + 2] = c;
            outBuffer[i * 8 + 3] = lfe;
            outBuffer[i * 8 + 4] = bc;
            outBuffer[i * 8 + 5] = bc;
            outBuffer[i * 8 + 6] = sl;
            outBuffer[i * 8 + 7] = sr;
        }
        else
        {
            // upmix N channels to N channels plus silence the empty channels
            float samples[inchannels];
            for (int j = 0; j < inchannels; ++j)
            {
                samples[j] = inBuffer[i * inchannels + j];
            }
            for (int j = 0; j < inchannels; ++j)
            {
                outBuffer[i * outchannels + j] = samples[j];
            }
            for (int j = inchannels; j < outchannels; ++j)
            {
                outBuffer[i * outchannels + j] = 0;
            }
        }
    }
}

@implementation DownmixProcessor

- (id) initWithInputFormat:(AudioStreamBasicDescription)inf andOutputFormat:(AudioStreamBasicDescription)outf {
    self = [super init];
    
    if (self) {
        if (inf.mFormatID != kAudioFormatLinearPCM ||
            (inf.mFormatFlags & kAudioFormatFlagsNativeFloatPacked) != kAudioFormatFlagsNativeFloatPacked ||
            inf.mBitsPerChannel != 32 ||
            inf.mBytesPerFrame != (4 * inf.mChannelsPerFrame) ||
            inf.mBytesPerPacket != inf.mFramesPerPacket * inf.mBytesPerFrame)
            return nil;
        
        if (outf.mFormatID != kAudioFormatLinearPCM ||
            (outf.mFormatFlags & kAudioFormatFlagsNativeFloatPacked) != kAudioFormatFlagsNativeFloatPacked ||
            outf.mBitsPerChannel != 32 ||
            outf.mBytesPerFrame != (4 * outf.mChannelsPerFrame) ||
            outf.mBytesPerPacket != outf.mFramesPerPacket * outf.mBytesPerFrame)
            return nil;
        
        inputFormat = inf;
        outputFormat = outf;

        [self setupVirt];

        [[NSUserDefaultsController sharedUserDefaultsController] addObserver:self forKeyPath:@"values.headphoneVirtualization" options:0 context:nil];
        [[NSUserDefaultsController sharedUserDefaultsController] addObserver:self forKeyPath:@"values.hrirPath" options:0 context:nil];
    }
    
    return self;
}

- (void) dealloc {
    [[NSUserDefaultsController sharedUserDefaultsController] removeObserver:self forKeyPath:@"values.headphoneVirtualization"];
    [[NSUserDefaultsController sharedUserDefaultsController] removeObserver:self forKeyPath:@"values.hrirPath"];
}

- (void) setupVirt {
    @synchronized(hFilter) {
        hFilter = nil;
    }
    
    BOOL hVirt = [[[NSUserDefaultsController sharedUserDefaultsController] defaults] boolForKey:@"headphoneVirtualization"];
    
    if (hVirt &&
        outputFormat.mChannelsPerFrame == 2 &&
        inputFormat.mChannelsPerFrame >= 1 &&
         inputFormat.mChannelsPerFrame <= 8) {
        NSString * userPreset = [[[NSUserDefaultsController sharedUserDefaultsController] defaults] stringForKey:@"hrirPath"];
        
        NSURL * presetUrl = nil;
        
        if (userPreset && ![userPreset isEqualToString:@""]) {
            presetUrl = [NSURL fileURLWithPath:userPreset];
            if (![HeadphoneFilter validateImpulseFile:presetUrl])
                presetUrl = nil;
        }
        
        if (!presetUrl) {
            presetUrl = [[NSBundle mainBundle] URLForResource:@"gsx" withExtension:@"wv"];
            if (![HeadphoneFilter validateImpulseFile:presetUrl])
                presetUrl = nil;
        }

        if (presetUrl) {
            @synchronized(hFilter) {
                hFilter = [[HeadphoneFilter alloc] initWithImpulseFile:presetUrl forSampleRate:outputFormat.mSampleRate withInputChannels:inputFormat.mChannelsPerFrame];
            }
        }
    }
}

- (void)observeValueForKeyPath:(NSString *)keyPath
                      ofObject:(id)object
                        change:(NSDictionary *)change
                       context:(void *)context
{
    DLog(@"SOMETHING CHANGED!");
    if ([keyPath isEqualToString:@"values.headphoneVirtualization"] ||
             [keyPath isEqualToString:@"values.hrirPath"]) {
        // Reset the converter, without rebuffering
        [self setupVirt];
    }
}

- (void) process:(const void *)inBuffer frameCount:(size_t)frames output:(void *)outBuffer {
    @synchronized (hFilter) {
        if ( hFilter ) {
            [hFilter process:(const float *) inBuffer sampleCount:frames toBuffer:(float *) outBuffer];
            return;
        }
    }
    
    if ( inputFormat.mChannelsPerFrame > 2 && outputFormat.mChannelsPerFrame == 2 )
    {
        downmix_to_stereo( (const float*) inBuffer, inputFormat.mChannelsPerFrame, (float*) outBuffer, frames );
    }
    else if ( inputFormat.mChannelsPerFrame > 1 && outputFormat.mChannelsPerFrame == 1 )
    {
        downmix_to_mono( (const float*) inBuffer, inputFormat.mChannelsPerFrame, (float*) outBuffer, frames );
    }
    else if ( inputFormat.mChannelsPerFrame < outputFormat.mChannelsPerFrame )
    {
        upmix( (const float*) inBuffer, inputFormat.mChannelsPerFrame, (float*) outBuffer, outputFormat.mChannelsPerFrame, frames );
    }
    else if ( inputFormat.mChannelsPerFrame == outputFormat.mChannelsPerFrame )
    {
        memcpy(outBuffer, inBuffer, frames * outputFormat.mBytesPerPacket);
    }
}

@end