use crate::audio::xa::{raw_audio::{CDAudioSamples, Orality}, Frequency, SampleDepth};
use cdtypes::types::sector::{Mode2Form2, XAADPCMBitsPerSample, XAADPCMSampleRate, XAADPCMSound};
use tool_helper::Error;

pub struct Encoder<'a> {
    left:              ChannelState,
    right:             ChannelState,
    source:            &'a[i16],
    frequency:         Frequency,
    sample_depth:      SampleDepth,
    orality:           Orality,
    samples_per_block: i32,
    sample_limit:      i32
}

impl<'a> Encoder<'a> {
    const BLOCKS_PER_SECTOR:usize    = 18;
    const XA_ADPCM_FILTER_COUNT: i32 = 4;
    const FILTER_K1: [i16; 5]        = [0, 60, 115,  98, 122];
    const FILTER_K2: [i16; 5]        = [0,  0, -52, -55, -60];

    pub fn new(cd_sample: &CDAudioSamples, frequency: Frequency, sample_depth: SampleDepth) -> Encoder {
        let orality                           = cd_sample.orality();
        let (samples_per_block, sample_limit) = Self::samples_per_block_and_limit(&cd_sample.samples(), sample_depth, orality);

        Encoder{left: ChannelState::default(), right: ChannelState::default(), source: &cd_sample.samples(), frequency, sample_depth, orality, samples_per_block, sample_limit}
    }

    pub fn encode_next_xa_sector(&mut self) -> Result<Option<Mode2Form2>, Error> {
        if self.source.is_empty() {
            return Ok(None);
        }

        let mut sector = self.create_new_sector();
        let mut dst    = &mut sector.data[0..];

        for _ in 0..Self::BLOCKS_PER_SECTOR {
            if self.source.len() < self.samples_per_block as usize {
                self.source = &self.source[self.source.len()..];
                break;
            }

            self.encode_xa(&self.source[0..], self.sample_limit, dst)?;
            
            self.sample_limit -= self.samples_per_block;
            self.source        = &self.source[self.samples_per_block as usize..];
            dst                = &mut dst[0x80..];
        }

        sector.finalize();
        Ok(Some(sector))
    }

    fn create_new_sector(&self) -> Mode2Form2 {
        let mut sector      = Mode2Form2::new();
        let     sub_mode    = &mut sector.sub_header.sub_mode;
        let     coding_info = &mut sector.sub_header.coding_info;

        sub_mode.set_real_time();
        
        coding_info.set_sound_type(match self.orality {
            Orality::Mono   => XAADPCMSound::Mono,
            Orality::Stereo => XAADPCMSound::Stereo
        });
        coding_info.set_sample_rate(match self.frequency {
            Frequency::Low  => XAADPCMSampleRate::Freq18900Hz,
            Frequency::High => XAADPCMSampleRate::Freq37800Hz,
        });
        coding_info.set_bits_per_sample(match self.sample_depth {
           SampleDepth::Normal => XAADPCMBitsPerSample::Normal,
           SampleDepth::High   => XAADPCMBitsPerSample::High, 
        });

        sector
    }

    fn encode_xa(&mut self, samples: &[i16], sample_limit: i32, data: &mut [u8]) -> Result<(), Error> {
        const SHIFT_RANGE_4BPS: i32 = 12;
        const SHIFT_RANGE_8BPS: i32 = 8;

        let channels = [&mut self.left, &mut self.right];
        match self.sample_depth {
            SampleDepth::Normal => {
                let (modulo, offset) = if self.orality == Orality::Stereo {(2, &STEREO_4BIT)} else {(1, &MONO_4BIT)};
                let (first_offset, second_offset) = offset;

                for (offset_idx, offset_set) in [first_offset, second_offset].iter().enumerate() {
                    for (idx, offset) in offset_set.iter().enumerate() {
                        let byte = Self::encode(channels[idx%modulo], &samples[offset.sample..], sample_limit + offset.sample_limit, offset.pitch, &mut data[offset.data..], offset.data_shift, offset.data_pitch, Self::XA_ADPCM_FILTER_COUNT, SHIFT_RANGE_4BPS)?;
                        data[idx + (offset_idx*8)]     = byte;
                        data[idx + 4 + (offset_idx*8)] = byte;
                    }
                }
            },
            SampleDepth::High   => {
                let (modulo, offset_set) = if self.orality == Orality::Stereo {(2, &STEREO_8BIT)} else {(1, &MONO_8BIT)};
                for (idx, offset) in offset_set.iter().enumerate() {
                    let byte = Self::encode(channels[idx%modulo], &samples[offset.sample..], sample_limit + offset.sample_limit, offset.pitch, &mut data[offset.data..], offset.data_shift, offset.data_pitch, Self::XA_ADPCM_FILTER_COUNT, SHIFT_RANGE_8BPS)?;    
                    data[idx]     = byte;
                    data[idx + 4] = byte;
                }
            }
        }

        Ok(())
    }

    fn encode(channel_state: &mut ChannelState, samples: &[i16], sample_limit: i32, pitch: i32, data: &mut [u8], data_shift: i32, data_pitch: i32, filter_count: i32, shift_range: i32) -> Result<u8, Error> {
        let mut best_mse          = 1i64 << 50i64;
        let mut best_filer        = 0;
        let mut best_sample_shift = 0;

        for filter in 0..filter_count {
            let true_min_shift = Self::find_min_shift(channel_state, samples, sample_limit, pitch, filter, shift_range)?;

            // Testing has shown that the optimal shift can be off the true minimum shift
            // by 1 in *either* direction.
            // This is NOT the case when dither is used.
            let min_shift = if true_min_shift - 1 < 0           {0}           else {true_min_shift - 1};
            let max_shift = if true_min_shift + 1 > shift_range {shift_range} else {true_min_shift + 1};

            for sample_shift in min_shift..=max_shift {
                let mut proposed = channel_state.clone();
                Self::attempt_encode(&mut proposed, samples, sample_limit, pitch, data, data_shift, data_pitch, filter, sample_shift, shift_range)?;
                if best_mse > proposed.mse {
                    best_mse          = proposed.mse;
                    best_filer        = filter;
                    best_sample_shift = sample_shift;
                }
            }
        }

        Self::attempt_encode(channel_state, samples, sample_limit, pitch, data, data_shift, data_pitch, best_filer, best_sample_shift, shift_range)
    }

    fn attempt_encode(out_channel_state: &mut ChannelState, samples: &[i16], sample_limit: i32, pitch: i32, data: &mut [u8], data_shift: i32, data_pitch: i32, filter: i32, sample_shift: i32, shift_range: i32) -> Result<u8, Error> {
        let sample_mask  = (0xFFFF >> shift_range) as u8;
        let nondata_mask = (!(sample_mask << data_shift)) as u8;

        let min_shift = sample_shift;
        let k1          = Self::FILTER_K1[filter as usize] as i32;
        let k2          = Self::FILTER_K2[filter as usize] as i32;

        let hdr = ((min_shift & 0x0F) | ((filter as i32) << 4)) as u8;

        out_channel_state.mse = 0;

        for i in 0..28 {
            let sample         = (if i >= sample_limit {0} else {samples[(i*pitch) as usize] as i32}) + out_channel_state.qerr;
            let previous_value = (k1*out_channel_state.prev1 + k2*out_channel_state.prev2 + (1 << 5)) >> 6;
            
            let mut sample_enc = sample - previous_value;
            sample_enc <<= min_shift;
            sample_enc +=  1 << (shift_range - 1);
            sample_enc >>= shift_range;

            if sample_enc < (std::i16::MIN as i32 >> shift_range) {sample_enc = std::i16::MIN as i32 >> shift_range}
            if sample_enc > (std::i16::MAX as i32 >> shift_range) {sample_enc = std::i16::MAX as i32 >> shift_range}
            sample_enc &= sample_mask as i32;

            let mut sample_dec = (((sample_enc & sample_mask as i32) << shift_range) as i16) as i32;
            sample_dec >>= min_shift;
            sample_dec += previous_value;
            if sample_dec > std::i16::MAX as i32 {sample_dec = std::i16::MAX as i32}
            if sample_dec < std::i16::MIN as i32 {sample_dec = std::i16::MIN as i32}

            let sample_error = sample_dec - sample;
            if sample_error >=  (1 << 30) || sample_error <= -(1 << 30) {
                return Err(Error::from_text(format!("Sample error exceeds 30bit: {}", sample_error)));
            }

            data[(i*data_pitch) as usize] = ((data[(i*data_pitch) as usize] & nondata_mask) as i32 | (sample_enc << data_shift)) as u8;

            out_channel_state.mse  += (sample_error as u64*sample_error as u64) as i64;
            out_channel_state.prev2 = out_channel_state.prev1;
            out_channel_state.prev1 = sample_dec;
        }

        Ok(hdr)
    }

    fn find_min_shift(channel_state: &ChannelState, samples: &[i16], sample_limit: i32, pitch: i32, filter: i32, shift_range: i32) -> Result<i32, Error> {
        /* 
            Assumption made:
            There is value in shifting right one step further to allow the nibbles to clip.
            However, given a possible shift value, there is no value in shifting one step less.

            Having said that, this is not a completely accurate model of the encoder,
            so maybe we will need to shift one step less.
        */

        let mut prev1       = channel_state.prev1;
        let mut prev2       = channel_state.prev2;
        let k1              = Self::FILTER_K1[filter as usize] as i32;
        let k2              = Self::FILTER_K2[filter as usize] as i32;
        let mut right_shift = 0;
        let mut s_min       = 0;
        let mut s_max       = 0;

        for i in 0..28 {
            let raw_sample  = if i >= sample_limit {0} else {samples[(i*pitch) as usize]} as i32;
            let prev_values = (k1*prev1 + k2*prev2 + (1 << 5)) >> 6;
            let sample      = raw_sample - prev_values;

            if sample < s_min {
                s_min = sample;
            }

            if sample > s_max {
                s_max = sample;
            }

            prev2 = prev1;
            prev1 = raw_sample;
        }

        while right_shift < shift_range && (s_max >> right_shift) > (std::i16::MAX as i32 >> shift_range) {
            right_shift += 1;
        }

        while right_shift < shift_range && (s_min >> right_shift) < (std::i16::MIN as i32 >> shift_range) {
            right_shift += 1;
        }

        let min_shift = shift_range - right_shift;
        if 0 <= min_shift && min_shift <= shift_range {
            Ok(min_shift)
        }

        else {
            Err(Error::from_text(format!("0 <= {} && {} <= {} was not satisfied with min_shift: {}", min_shift, min_shift, shift_range, min_shift)))
        }
    }

    fn samples_per_block_and_limit(input: &[i16], sample_depth: SampleDepth, orality: Orality) -> (i32, i32) {        
        let samples_per_block = match sample_depth {
            SampleDepth::Normal => 224,
            SampleDepth::High   => 112,
        };
    
        let sample_limit = match orality {
            Orality::Stereo => input.len()*2,
            Orality::Mono   => input.len(),
        };
    
        (samples_per_block, sample_limit as i32)
    }
}

#[derive(Clone)]
struct ChannelState {
    qerr:  i32, // quanitisation error
    mse:   i64, // mean square error
    prev1: i32, 
    prev2: i32
}

impl std::default::Default for ChannelState {
    fn default() -> Self { 
        ChannelState{qerr: 0, mse: 0, prev1: 0, prev2: 0}
    }
}

struct EncodingOffsets {
    sample: usize,
    sample_limit: i32,
    pitch: i32,
    data: usize,
    data_shift: i32,
    data_pitch: i32,
}

const STEREO_4BIT: ([EncodingOffsets; 4], [EncodingOffsets; 4]) = (
    [
        EncodingOffsets{sample: 0,      sample_limit:   0, pitch: 2, data: 0x10, data_shift: 0, data_pitch: 4},
        EncodingOffsets{sample: 1,      sample_limit:   0, pitch: 2, data: 0x10, data_shift: 4, data_pitch: 4},
        EncodingOffsets{sample: 56,     sample_limit: -28, pitch: 2, data: 0x11, data_shift: 0, data_pitch: 4},
        EncodingOffsets{sample: 56 + 1, sample_limit: -28, pitch: 2, data: 0x11, data_shift: 4, data_pitch: 4},
    ],
    [
        EncodingOffsets{sample: 56*2,     sample_limit: -28*2, pitch: 2, data: 0x12, data_shift: 0, data_pitch: 4},
        EncodingOffsets{sample: 56*2 + 1, sample_limit: -28*2, pitch: 2, data: 0x12, data_shift: 4, data_pitch: 4},
        EncodingOffsets{sample: 56*3,     sample_limit: -28*3, pitch: 2, data: 0x13, data_shift: 0, data_pitch: 4},
        EncodingOffsets{sample: 56*3 + 1, sample_limit: -28*3, pitch: 2, data: 0x13, data_shift: 4, data_pitch: 4}
    ]
);
const MONO_4BIT: ([EncodingOffsets; 4], [EncodingOffsets; 4]) = (
    [
        EncodingOffsets{sample: 0,    sample_limit:     0, pitch: 1, data: 0x10, data_shift: 0, data_pitch: 4},
        EncodingOffsets{sample: 28,   sample_limit:   -28, pitch: 1, data: 0x10, data_shift: 4, data_pitch: 4},
        EncodingOffsets{sample: 28*2, sample_limit: -28*2, pitch: 1, data: 0x11, data_shift: 0, data_pitch: 4},
        EncodingOffsets{sample: 28*3, sample_limit: -28*3, pitch: 1, data: 0x11, data_shift: 4, data_pitch: 4},
    ],
    [
        EncodingOffsets{sample: 28*4, sample_limit: -28*4, pitch: 1, data: 0x12, data_shift: 0, data_pitch: 4},
        EncodingOffsets{sample: 28*5, sample_limit: -28*5, pitch: 1, data: 0x12, data_shift: 4, data_pitch: 4},
        EncodingOffsets{sample: 28*6, sample_limit: -28*6, pitch: 1, data: 0x13, data_shift: 0, data_pitch: 4},
        EncodingOffsets{sample: 28*7, sample_limit: -28*7, pitch: 1, data: 0x13, data_shift: 4, data_pitch: 4}
    ]
);

const STEREO_8BIT: [EncodingOffsets;4] = [
    EncodingOffsets{sample:  0,     sample_limit:   0, pitch: 2, data: 0x10, data_shift: 0, data_pitch: 4},
    EncodingOffsets{sample:  1,     sample_limit:   0, pitch: 2, data: 0x11, data_shift: 0, data_pitch: 4},
    EncodingOffsets{sample: 56,     sample_limit: -28, pitch: 2, data: 0x12, data_shift: 0, data_pitch: 4},
    EncodingOffsets{sample: 56 + 1, sample_limit: -28, pitch: 2, data: 0x13, data_shift: 0, data_pitch: 4},
];
const MONO_8BIT:   [EncodingOffsets;4] = [
    EncodingOffsets{sample:    0, sample_limit:     0, pitch: 1, data: 0x10, data_shift: 0, data_pitch: 4},
    EncodingOffsets{sample:   28, sample_limit:   -28, pitch: 1, data: 0x11, data_shift: 0, data_pitch: 4},
    EncodingOffsets{sample: 28*2, sample_limit: -28*2, pitch: 1, data: 0x12, data_shift: 0, data_pitch: 4},
    EncodingOffsets{sample: 28*3, sample_limit: -28*3, pitch: 1, data: 0x13, data_shift: 0, data_pitch: 4},
];