from reg_map import *

class Reference_model:
    def __init__(
        self,
        dut,
        pulse_width,
        pulse_period,
        pulse_num,
        pulse_height,
        adc_period,
        window_size,
        DAC_DATA_WIDTH,
        ADC_DATA_WIDTH,
        PACK_FACTOR,
        PROCESS_MODE,
        ZERO_LEVEL,
        ACCUM_WIDTH,
        N_MAX,
        PACKET_SIZE,
        RD_FIFO_WIDTH
    ):

        self.dut = dut

        # configuration
        self.pulse_width = pulse_width
        self.pulse_period = pulse_period
        self.pulse_num = pulse_num
        self.pulse_height = pulse_height
        self.adc_period = adc_period
        self.window_size = window_size

        # parameters
        self.DAC_DATA_WIDTH = DAC_DATA_WIDTH
        self.ADC_DATA_WIDTH = ADC_DATA_WIDTH
        self.PACK_FACTOR = PACK_FACTOR
        self.PROCESS_MODE = PROCESS_MODE
        self.ZERO_LEVEL = ZERO_LEVEL
        self.ACCUM_WIDTH = ACCUM_WIDTH
        self.N_MAX = N_MAX
        self.PACKET_SIZE = PACKET_SIZE
        self.RD_FIFO_WIDTH = RD_FIFO_WIDTH

        # intermediate data
        self.samples = []
        self.window_data = []
        self.accum_data = []


    def gen_input_samples (self):
        self.samples = []
        DAC_ZERO = self.ZERO_LEVEL
        ADC_ZERO = 1 << (self.ADC_DATA_WIDTH - 1)

        DAC_RANGE = 5.0
        ADC_RANGE = 1.0

        GROUND_BIAS = 0.0
        ADC_GAIN = 0.2

        DAC_STEP = (2 * DAC_RANGE) / ((1 << self.DAC_DATA_WIDTH) - 1)
        ADC_STEP = (2 * ADC_RANGE) / ((1 << self.ADC_DATA_WIDTH) - 1)

        ADC_MAX = (1 << self.ADC_DATA_WIDTH) - 1

        for _ in range(self.pulse_num):
            sample_time = 0

            while sample_time < self.pulse_period:
                if sample_time < self.pulse_width:
                    dac_code = self.pulse_height
                else:
                    dac_code = DAC_ZERO

                voltage = (dac_code - DAC_ZERO) * DAC_STEP
                voltage = (voltage - GROUND_BIAS) * ADC_GAIN

                if voltage <= -ADC_RANGE:
                    adc_code = 0
                elif voltage >= ADC_RANGE:
                    adc_code = ADC_MAX
                else:
                    adc_code = int(round(voltage / ADC_STEP + ADC_ZERO))

                out_of_range = abs(voltage) >= ADC_RANGE #(adc_code == 0 or adc_code == ADC_MAX)

                if self.PROCESS_MODE:
                    msb = (adc_code >> (self.ADC_DATA_WIDTH - 1)) & 1

                    if out_of_range:
                        if msb:
                            sample = (1 << self.ADC_DATA_WIDTH) - 1
                        else:
                            sample = 0

                    else:
                        sample = (((~msb) & 1) << (self.ADC_DATA_WIDTH - 1)) | (adc_code & ((1 << (self.ADC_DATA_WIDTH - 1)) - 1))

                else:
                    if out_of_range:

                        if adc_code & (1 << (self.ADC_DATA_WIDTH - 1)):
                            sample = (1 << self.ADC_DATA_WIDTH) - 1
                        else:
                            sample = 0

                    else:
                        sample = adc_code

                self.samples.append(sample)
                sample_time += self.adc_period


        
    def apply_window (self):
        self.window_data = []
        if self.window_size < 1:
            raise ValueError(f"window_size must be >= 1, got {self.window_size}" )

        if len(self.samples) == 0:
            raise ValueError( "samples[] is empty. Call gen_input_samples() first." )

        if len(self.samples) % self.window_size != 0:
            raise ValueError( f"Number of samples ({len(self.samples)}) " f"is not divisible by window_size ({self.window_size})" )
        
        accum = 0
        cnt = 0

        for sample in self.samples:
            accum += sample
            cnt += 1
            if cnt == self.window_size:
                self.window_data.append(accum)
                accum = 0
                cnt = 0



    def accumulate (self):
        self.accum_data = []
        windows_per_pulse = len(self.window_data) // self.pulse_num

        if windows_per_pulse == 0:
            raise ValueError("No windows were generated.")

        if len(self.window_data) == 0:
            raise ValueError( "window_data is empty. Call apply_window() first.")

        if len(self.window_data) % self.pulse_num != 0:
            raise ValueError( "window_data length is not divisible by pulse_num")
        
        for window in range(windows_per_pulse):
            accum = 0

            for pulse in range(self.pulse_num):

                index = pulse * windows_per_pulse + window
                accum += self.window_data[index]

            self.accum_data.append(accum)
            if len(self.accum_data) == windows_per_pulse:
                print( f"[ReferenceModel] Accumulation complete: "
            f"{self.pulse_num} pulses -> "
            f"{windows_per_pulse} windows per pulse -> "
            f"{len(self.accum_data)} output values.")
    

    def run(self):
        self.gen_input_samples()
        self.apply_window()
        self.accumulate()
        return self.accum_data