# general imports
import sys
from inspect import stack   # get name of caller function
from typing import Callable
from math import log2, ceil
from enum import Enum
import random

#hdl specific imports
from amaranth import *
from amaranth import Elaboratable
from amaranth.back import verilog, cxxrtl, rtlil
from amaranth.sim import Settle, Delay, Simulator

#custom imports 
from hdl.config import *

# configuration
random.seed(5498)

__all__ = ['cmd', 'e2s', 'sim', 'step', 'eval', 'rand_bits', 'rand_bits_extremes', 'rand_bits_mix']


def cmd(hdl):
    '''
    Very simple command line interface
    The elaboratable class must have a ports attribute that is a dict of in and out ports {'in': [Signals()], 'out': [Signals()]}
    '''

    if len(sys.argv) <= 1:
        print('Usage: v|cc v = generate verilog, cc = generate cxxrtl')
        exit()

    if sys.argv[1] == "sim":
        # tb(sys.argv[0].replace('.py', '.vcd'))
        # exit()
        assert "sim option deprecated, use pytest command instead"

    if sys.argv[1] == "v":
        out = verilog.convert(hdl, ports=hdl.ports['in'] + hdl.ports['out'])
        with open(os.path.join(VERILOG_DIR, os.path.basename(sys.argv[0]).replace('.py', '.v')), 'w') as f:
            f.write(out)
    
    elif sys.argv[1] == "cc":
        out = cxxrtl.convert(hdl, ports=hdl.ports['in'] + hdl.ports['out'])
        with open(os.path.join(CXXRTL_DIR, os.path.basename(sys.argv[0]).replace('.py', '.cc')), 'w') as f:
            f.write(out)

    elif sys.argv[1] == "rt":
        out = rtlil.convert(hdl, ports=hdl.ports['in'] + hdl.ports['out'])
        with open(os.path.join(RTLIL_DIR, os.path.basename(sys.argv[0]).replace('.py', '.il')), 'w') as f:
            f.write(out)

def e2s(e: Enum):
    '''
    Get signal length from enum, returns ceil(log2(len(e)))) 
    '''
    return ceil(log2(len(e)))

def sim(dut:Elaboratable, proc: Callable, sync=True):
    sim = Simulator(dut)

    if sync:
        sim.add_clock(1e-9)
        sim.add_sync_process(proc)
    else:
        sim.add_process(proc)
    
    with sim.write_vcd(os.path.join(VCD_DIR, stack()[1].function + '.vcd')):    # get name of caller function
        sim.run()

def step(cycles=1):
    for _ in range(cycles):
        yield Settle()  # settle comb logic before clock
        yield           # clock edge
        yield Settle()  # settle comb logic after clock

def eval():
    print(DeprecationWarning('eval() is deprecatated and should not be used'))
    yield Settle()
    yield Delay(1e-6)
    

# bits can be integer or Signal
def rand_bits(bits, sus=None, low=None, high=None):
    '''
    Takes number of bits or a Signal and returns a random number
    if low and high are not specified, the range is 0 to 2**bits - 1
    If low and high are specified the number of bits is ignored
    sus (Signed UnSigned) can be 's' or 'u' to force signed or unsigned, default is random
    sus has no effect if low and high are specified
    '''

    if isinstance(bits, Signal):
        bits = bits.width

    if (low == high == None):
        if sus == 's':
            low = -2**(bits-1)
            high = 2**(bits-1) - 1
        elif sus == 'u':
            low = 0
            high = 2**bits - 1
        else:
            if random.random() < 0.5:
                low = 0
                high = 2**bits - 1
            else:
                low = -2**(bits-1)
                high = 2**(bits-1) - 1
    else:
        assert low < high, "low must be less than high"
        assert low != None, "low must be defined if high is defined"
        assert high != None, "high must be defined if low is defined"

    return random.randint(low, high)

# bits can be integer or Signal
def rand_bits_extremes(bits, sus=None):
    '''
    Takes number of bits or a Signal and returns a random number with the extremes of the range
    sus (Signed UnSigned) can be 's' or 'u' to force signed or unsigned, default is random
    '''

    if isinstance(bits, Signal):
        bits = bits.width

    if sus == 's':
        low = -2**(bits-1)
        high = 2**(bits-1) - 1
    elif sus == 'u':
        low = 0
        high = 2**bits - 1
    else:
        if random.random() < 0.5:
            low = 0
            high = 2**bits - 1
        else:
            low = -2**(bits-1)
            high = 2**(bits-1) - 1
        
    choices = [low, high] if low == 0 else [low, 0, high]   # if low is 0, don't include it as a choice
    return random.choice(choices)


def rand_bits_mix(bits: int|Signal, low:int=None, high:int=None, sus:"None, 'u', 's'"=None, extrem_prob:float=0.25):
    '''
    Take number of bits or a Signal
    sus determines if the number is signed or unsigned, or both ('s', 'u', or None)
    Returns a random number with increased likelihood of extremes
    If low and high are specified the number of bits is ignored
    The likelyhood of extremes is determined by extrem_prob, default is 0.25
    '''
    if random.random() <= extrem_prob:
        return rand_bits_extremes(bits, sus=sus)
    else:
        return rand_bits(bits, sus=sus, low=low, high=high)