from migen import *

from litex.soc.interconnect.wishbone import *

from math import log2, ceil
from typing import List

"""
Random implementation notes:

- Circular buffers can keep overwriting. We only need a setting to say how many samples to save after
  trigger occurs.
- Data valid from samplers to FIFOs can simply be gated via the enable signal. Everything can just run
  all the time to keep things simple
- can we correct clock skew on the sample clock via Lattice primitives? I think it's possible. I doubt it
  matters. Would need significant calibration effort to even have it be accurate.
- Trigger system should wait a couple clocks after trigger acquired to disable FIFOs, just in case the
  CDC sync happens a bit late for some ADC channels

Configurable parameters:
- trigger_run_len: number of samples to acquire after triggered sample (can technically be arbitrarily
large, circular buffer handles data loss, should be larger than trigger_thresh_time to make sure buffers
don't get weird)
- trigger_thresh_value: minimum peak to peak value to consider triggered
- trigger_thresh_time: minimum num samples that peak must be above threshold to count as a trigger
                        (trigger sample number is the first sample above the threshold value) (must be >= 1)
- trigger_decay_value: decay value to subtract from peak values to potentially reduce false triggers
- trigger_decay_period: number of samples per decay application


Implementation of trigger (psuedocode), happens every sample update:

if triggered:
    if num_samples + 1 >= trigger_run:
        disable_trigger()
        return

    num_samples += 1
    return

if sample > max:
    max = sample
elif sample < min:
    min = sample

if (max - min) > trigger_thresh_value:
    if triggered_for + 1 >= trigger_thresh_time:
        triggered = True
        num_samples = 0
        return

    triggered_for += 1
    decay_wait = 0
else:
    triggered_for = 0
    decay_wait += 1

    if trigger_decay_period == 0 or decay_wait == trigger_thresh_time:
        decay_wait = 0

        if (max - trigger_decay_value) > (min + trigger_decay_value):
            max -= trigger_decay_value
            min += trigger_decay_value
"""

class CircularBuffer(Module):
    """
    Circular buffer implementation that allows users to read the entire data.

    Assumptions:
    - Reading values while writes are ocurring does not need to have well-defined behaviour

    Implementation is largely based on Migen SyncFIFO, just tweaked to operate how I want
    """

    def __init__(self, width: int, depth: int, with_wb = True) -> None:
        storage = Memory(width=width, depth=depth)
        self.specials += storage

        ptr_width = ceil(log2(depth))

        # External Signals
        self.len = Signal(ptr_width) # Amount of valid data in the buffer
        self.clear = Signal()                # Strobe to clear memory
        self.rd_addr = Signal(ptr_width)
        self.rd_data = Signal(width)

        self.wr_data = Signal(width)
        self.wr_ready = Signal() # Output, signals buffer is ready to be written to
        self.wr_valid = Signal() # Input, high when data is present to be written

        wr_ptr = Signal(ptr_width)
        rd_ptr = Signal(ptr_width)
        empty = Signal(reset=1) # Extra signal to distinguish between full and empty condition

        # Hook write input signals to memory
        wr_port = storage.get_port(write_capable=True)
        # Always ready to write data into memory, so hook these signals straight in
        self.comb += [
            wr_port.adr.eq(wr_ptr),
            wr_port.dat_w.eq(self.wr_data),
            wr_port.we.eq(self.wr_valid),
            self.wr_ready.eq(1),            # We are always ready to write data in
        ]

        # Advance write (and potentially read)
        self.sync += [
            If(self.wr_valid,
                # We aren't empty anymore, and we won't be until we are cleared
                empty.eq(0),

                # Advance write pointer
                If(wr_ptr < (depth - 1),
                   wr_ptr.eq(wr_ptr + 1))
                .Else(wr_ptr.eq(0)),

                # Advance read pointer if we are full (e.g. overwrite old data)
                If(~empty & (wr_ptr == rd_ptr),
                    If(rd_ptr < (depth - 1),
                       rd_ptr.eq(rd_ptr + 1))
                    .Else(rd_ptr.eq(0))
                )
            )
        ]

        # TODO should I actually set async_read?
        rd_port = storage.get_port(async_read=True)
        # Set read addr so 0 starts at rd_ptr and wraps around, and connect read data up
        self.comb += [
            If(self.rd_addr + rd_ptr < depth,
               rd_port.adr.eq(self.rd_addr + rd_ptr))
            .Else(
                rd_port.adr.eq(self.rd_addr - (depth - rd_ptr))
            ),
            self.rd_data.eq(rd_port.dat_r),
        ]

        # Export the length present
        self.comb += [
            If(empty, self.len.eq(0))
            .Else(
                If(wr_ptr > rd_ptr,
                    self.len.eq(wr_ptr - rd_ptr))
                .Elif(wr_ptr != rd_ptr,
                    self.len.eq(depth - (rd_ptr - wr_ptr)))
                .Else(
                    self.len.eq(depth)
                )
            ),
        ]

        # "Clear" out memory if clear is strobed
        # NOTE really clear should be hooked into reset, but I'm not clear on how to do that.
        # Technically there's some glitches that can happen here if we write data while clear
        # is asserted, but that shouldn't happen and it's fine if it does tbh.
        self.sync += If(self.clear,
                        wr_ptr.eq(0), rd_ptr.eq(0), empty.eq(1))

        # Add wishbone bus to access data
        if with_wb:
            self.bus = Interface(data_width=32, adr_width=ceil(log2(depth)))

            self.comb += self.rd_addr.eq(self.bus.adr)
            self.sync += [
                self.bus.ack.eq(0),
                self.bus.dat_r.eq(0),
                If(~self.bus.we & self.bus.cyc & self.bus.stb,
                   self.bus.ack.eq(1), self.bus.dat_r.eq(self.rd_data)),
            ]


from migen.genlib.cdc import PulseSynchronizer


class Sampler(Module):
    def __init__(self, adc_pins: Record, sampler_clock: Signal):
        # TODO remove bus
        self.bus = Interface(data_width=32, adr_width=11)

        # self.clock_domains.foo = ClockDomain() is how to add a new clock domain, accessible at self.foo
        # Connect sampler clock domain
        self.clock_domains.sample_clock = ClockDomain("sample_clock")
        self.comb += self.sample_clock.clk.eq(sampler_clock)

        # Hook up ADC REFCLK to sample_clock
        self.comb += adc_pins.refclk.eq(sampler_clock)

        # We can synchronize to the sampler clock, whenever it goes high we can
        # strobe a single valid signal
        synchronizer = PulseSynchronizer("sample_clock", "sys")
        self.submodules += synchronizer

        self.valid = Signal()
        self.data = Signal(10)

        self.comb += [
            synchronizer.i.eq(self.sample_clock.clk),
            self.valid.eq(synchronizer.o),
            self.data.eq(adc_pins.data),
        ]

        # Set config pins to constant values
        self.comb += adc_pins.oen_b.eq(0)       # Data pins enable
        self.comb += adc_pins.standby.eq(0)     # Sampling standby
        self.comb += adc_pins.dfs.eq(0)         # DFS (raw or two's complement)
        # The only remaining pin, OTR, is an out of range status indicator

        # Read directly from the data pins into the wishbone bus for now, just for bringup
        self.sync += If(self.valid, self.bus.dat_r.eq(adc_pins.data))
        self.sync += self.bus.ack.eq(0)
        self.sync += If(self.bus.cyc & self.bus.stb, self.bus.ack.eq(1))


class PeakDetector(Module):
    """
    Module to detect when peak to peak voltage is high enough to consider incoming
    data to be a valid ping. Configuration is provided by setting the configuration
    attributes. Do not change these settings while detector is running.

    Attributes
    ----------
    data: (input)
        Data signal to use for detection

    data_valid: (input)
        Strobed signal that indicates value on `data` is valid to be read

    enable: (input)
        Enables running peak detection. De-asserting this will clear all state variables

    triggered: (output)
        Signal that indicates peak has been triggered. Only cleared once enable is de-asserted again

    Configuration Attributes
    ------------------------
    thresh_value:
        Minimum peak to peak value considered triggered

    thresh_time:
        Number of consecutive samples above threshold required to consider triggered

    decay_value:
        Decay value to subtract from peak values to prevent false triggers

    decay_period:
        Number of samples between each application of decay
    """

    def __init__(self, data_width: int):
        # Create all state signals
        min_val = Signal(data_width)
        max_val = Signal(data_width)
        diff = Signal(data_width)
        triggered_time = Signal(32)
        decay_counter = Signal(32)

        # Control signals
        self.data = Signal(data_width)
        self.data_valid = Signal()
        self.enable = Signal()
        self.triggered = Signal()

        # Configuration Parameters
        self.thresh_value = Signal(data_width)
        self.thresh_time = Signal(32)
        self.decay_value = Signal(data_width)
        self.decay_period = Signal(32)

        self.sync += If(~self.enable,
            # Reset halfway. ADCs are 0-2V, and everything should be centered at 1V, so this is approximating the initial value
            min_val.eq(int(2**data_width /2)),
            max_val.eq(int(2**data_width /2)),
            self.triggered.eq(0),
            decay_counter.eq(0),
            triggered_time.eq(0),
        )

        # Constantly updating diff to simplify some statements
        self.comb += diff.eq(max_val - min_val)

        self.sync += If(self.enable & self.data_valid,
            # Update maximum value
            If(self.data > max_val, max_val.eq(self.data)),
            # Update minimum value
            If(self.data < min_val, min_val.eq(self.data)),
            If(diff > self.thresh_value,
               # We have met the threshold for triggering, start counting
               triggered_time.eq(triggered_time + 1),
               decay_counter.eq(0),

               # We have triggered, so we can set the output. After this point,
               # nothing we do matters until enable is de-asserted and we reset
               # triggered.
               If(triggered_time + 1 >= self.thresh_time, self.triggered.eq(1)))
            .Else(
                # We have not met the threshold, reset timer and handle decay
                triggered_time.eq(0),
                decay_counter.eq(decay_counter + 1),

                # Decay threshold has been reached, apply decay to peaks
                If(decay_counter >= self.decay_period,
                   decay_counter.eq(0),

                   # Only apply decay if the values would not overlap
                   If(diff >= (self.decay_value << 1),
                      max_val.eq(max_val - self.decay_value),
                      min_val.eq(min_val + self.decay_value)))
            )
        )


class SamplerController(Module):
    """
    Sampler control

    Attributes
    ----------
    bus:
        Slave wishbone bus to be connected to a higher-level bus. Has an address width set according to
        the provided buffer length.

    buffers:
        List of FIFO buffer objects used to store sample data.

    samplers:
        List of sampler objects provided by user.

    Registers
    --------
    0x00: Control Register (RW)
        Bit 0 - Begin capture. Resets all FIFOs and starts the peak detector

    0x01: Status Register (RO)
        Bit 0 - Capture complete. Set by peak detection block and cleared by software or when 

    0x02: trigger_run_len (RW)
        Number of samples to acquire after triggering sample.

    0x03: thresh_value (RW)
        Minimum peak to peak value considered triggered

    0x04: thresh_time (RW)
        Number of consecutive samples above threshold required to consider triggered

    0x05: decay_value (RW)
        Decay value to subtract from peak values to prevent false triggers

    0x06: decay_period (RW)
        Number of samples between each application of decay

    0x1xx: BUFFER_LEN_X (RO)
        Lenght of data in buffer, up to the number of samplers provided.

    """
    def __init__(self, samplers: List[Sampler], buffer_len):

        self.samplers = samplers
        num_channels = len(samplers)

        # Enables reading in samples
        sample_enable = Signal()
        # Pull in only one CDC sync signal
        sample_ready = self.samplers[0].valid

        # Generate buffers for each sampler
        self.buffers = [CircularBuffer(9, buffer_len) for _ in range(num_channels)]

        # Connect each buffer to each sampler
        for buffer, sampler in zip(self.buffers, self.samplers):
            self.comb += [
                # Connect only top 9 bits to memory
                buffer.wr_data.eq(sampler.data[1:]),
                # Writes enter FIFO only when enabled and every clock cycle
                buffer.wr_valid.eq(sample_enable & sample_ready),
            ]


        # Each sampler gets some chunk of memory at least large enough to fit
        # all of it's data, so use that as a consistent offset
        sample_mem_addr_width = ceil(log2(buffer_len))
        # 1 control block + number of channels used = control bits
        control_block_addr_width = ceil(log2(num_channels + 1))

        # Bus address width
        addr_width = control_block_addr_width + sample_mem_addr_width

        # "Master" bus
        self.bus = Interface(data_width=32, addr_width=addr_width)

        # Wishbone bus used for mapping control registers
        self.control_regs_bus = Interface(data_width=32, addr_width=sample_mem_addr_width)

        slaves = []
        slaves.append((lambda adr: adr[sample_mem_addr_width:] == 0, self.control_regs_bus))

        for i, buffer in enumerate(self.buffers):
            # Connect subordinate buses of buffers to decoder
            slaves.append((lambda adr: adr[sample_mem_addr_width:] == i + 1, buffer.bus))

            adr = (i + 1) << sample_mem_addr_width
            print(f"Sampler {i} available at 0x{adr:08x}")

        self.decoder = Decoder(self.bus, slaves)
        # TODO how to submodule
        self.submodules.decoder = self.decoder

        self.peak_detector = PeakDetector(10)
        self.comb += [
            # Simply enable whenever we start capturing
            self.peak_detector.enable.eq(sample_enable),
            # Connect to the first ADC
            self.peak_detector.data.eq(self.samplers[0].data),
            # Use the same criteria as the fifo buffer
            self.peak_detector.data_valid.eq(sample_enable & sample_ready),
        ]

        #### Control register logic

        # Storage
        control_register = Signal(32)
        status_register = Signal(32)
        trigger_run_len = Signal(32)

        def rw_register(storage: Signal, *, read: bool = True, write: bool = True):
            if read:
                read = self.control_regs_bus.dat_r.eq(storage)
            else:
                read = self.control_regs_bus.ack.eq(0)

            if write:
                write = storage.eq(self.control_regs_bus.dat_w)
            else:
                write = self.control_regs_bus.ack.eq(0)

            return If(self.control_regs_bus.we, write).Else(read)

        # Handle explicit config registers
        cases = {
            0: rw_register(control_register),
            1: rw_register(status_register, write=False),
            2: rw_register(trigger_run_len),
            3: rw_register(self.peak_detector.thresh_value),
            4: rw_register(self.peak_detector.thresh_time),
            5: rw_register(self.peak_detector.decay_value),
            6: rw_register(self.peak_detector.decay_period),

            "default": rw_register(None, read=False, write=False)
        }

        # Handle length values for each sample buffer
        for i, buffer in enumerate(self.buffers):
            cases.update({0x100 + i: rw_register(buffer.len, write=False)})

        # Connect up control registers bus
        self.sync += [
            self.control_regs_bus.ack.eq(0),
            If(self.control_regs_bus.cyc & self.control_regs_bus.stb,
                        self.control_regs_bus.ack.eq(1),
                        Case(self.control_regs_bus.adr, cases)),
        ]

        # Handle the control logic
        post_trigger_count = Signal(32)
        self.sync += [
            # Reset state whenever sampling is disabled
            If(~sample_enable, post_trigger_count.eq(0)),

            # Reset triggering status if we have started sampling
            # (peak_detector.triggered resets if sample_enable is de-asserted, so
            # this is a reliable reset mechanism)
            If(sample_enable & ~self.peak_detector.triggered,
               status_register[0].eq(0)),

            # Keep sampling past the trigger for the configured number of samples
            If(self.peak_detector.triggered & sample_enable & sample_ready,
               post_trigger_count.eq(post_trigger_count + 1),

               # We have sampled enough, update status and stop sampling
               If(post_trigger_count + 1 >= trigger_run_len,
                  status_register[0].eq(1),
                  control_register[0].eq(0))),
        ]

        # Update register storage
        self.comb += [
            sample_enable.eq(control_register[0]),
        ]


def fifo_testbench():
    dut = CircularBuffer(9, 24)
    def test_fn():
        assert (yield dut.len) == 0
        assert (yield dut.wr_ready) == 1

        # Clock some data in, check len
        data = [0xDE, 0xAD, 0xBE, 0xEF]
        for b in data:
            (yield dut.wr_data.eq(b))
            (yield dut.wr_valid.eq(1))
            yield

        # Stop clocking data in
        (yield dut.wr_valid.eq(0))
        # Tick again because setting a value waits until the next clock...
        yield

        fifo_len = (yield dut.len)
        assert fifo_len == 4, f"len should be 4, is {fifo_len}"

        # Reset
        (yield dut.clear.eq(1))
        yield
        (yield dut.clear.eq(0))
        yield

        # Len should be cleared
        assert (yield dut.len) == 0

        # Clock more data in than capacity, check that we can read out
        # the expected data
        data = [r for r in range(32)] # Yes yes I could use a generator but I want to slice it later
        for b in data:
            (yield dut.wr_data.eq(b))
            (yield dut.wr_valid.eq(1))
            yield

        # One more clock
        (yield dut.wr_valid.eq(0))
        yield

        data_len = (yield dut.len)
        assert data_len == 24, f"len should be 24, is {data_len}"
        out_data = []
        for i in range(24):
            (yield dut.rd_addr.eq(i))
            yield

            out_data.append((yield dut.rd_data))

            assert out_data[i] == data[i + 8], f"Data mismatch at index {i}, should be {data[i+8]}, is {out_data[i]}"

        # At this point, everything seems to be good, so I'm leaving more exhaustive testing

    run_simulation(dut, test_fn())


def write_wishbone(bus, address, value):
    # Set up bus
    (yield bus.adr.eq(address))
    (yield bus.dat_w.eq(value))
    (yield bus.stb.eq(1))
    (yield bus.cyc.eq(1))
    (yield bus.we.eq(1))
    yield

    cycles = 0
    while True:
        cycles += 1
        assert cycles < 5, "Write fail"


        if (yield bus.ack) == 1:
            # We received a response, clear out bus status and exit
            (yield bus.stb.eq(0))
            (yield bus.cyc.eq(0))
            yield

            break
        else:
            # Tick until we receive an ACK
            yield


def read_wishbone(bus, address,):
    """Sets up a read transaction. Due to limitations of the simulation method, you have to read
    from dat_r, and also tick immediately after calling"""
    # Set up bus
    (yield bus.adr.eq(address))
    (yield bus.stb.eq(1))
    (yield bus.cyc.eq(1))
    (yield bus.we.eq(0))
    yield

    cycles = 0
    while True:
        cycles += 1
        assert cycles < 5, "Write fail"
        if (yield bus.ack) == 1:
            # We received a response, clear out bus status and exit
            (yield bus.stb.eq(0))
            (yield bus.cyc.eq(0))

            break
        else:
            # Tick until we receive an ACK
            yield

class MockSampler(Module):
    """
    Attributes
    ----------
    All Sampler attributes by default, plus the following:

    index:
        Index of data to use from provided data
    """
    def __init__(self, data: List[int]):
        memory = Memory(width=10, depth=len(data), init=data)

        self.index = Signal(ceil(log2(len(data))))
        self.data = Signal(10)
        self.valid = Signal()

        read_port = memory.get_port(async_read=True)
        self.comb += [
            read_port.adr.eq(self.index),
            self.data.eq(read_port.dat_r),
        ]

class TestSoC(Module):
    def __init__(self, data):
        sampler = MockSampler(data)
        self.submodules.sampler = sampler
        # TODO multiple mock samplers to test that functionality
        self.controller = SamplerController([MockSampler(data)], 1024)
        self.submodules.controller = self.controller
        self.bus = self.controller.bus


def controller_test_bus_access():
    dut = TestSoC([2, 3, 4, 5])
    def test_fn():
        yield from write_wishbone(dut.bus, 2, 0xDEADBEEF)
        yield from read_wishbone(dut.bus, 2)
        assert (yield dut.bus.dat_r) == 0xDEADBEEF, "Read failed!"

    # TODO test writing to RO register fails

    run_simulation(dut, test_fn(), vcd_name="test_bus_access.vcd")

# TODO test a couple variations on waveforms:
# Just a clean waveform, should pass normally
# Clean waveform w/ some decay
# Some waveform that decay could make not trigger (i.e. a big spike)
# Clean waveform under threshold
# Test that decay operates normally and settles back down to center value

if __name__ == "__main__":
    import argparse

    args = argparse.ArgumentParser()
    args.add_argument("--fifo", action="store_true", help="Run FIFO tests")
    args.add_argument("--controller", action="store_true", help="Run sampler tests")
    args = args.parse_args()

    if args.fifo:
        fifo_testbench()

    if args.controller:
        controller_test_bus_access()
