import enum
import wave
import math
import random
import dataclasses
from dataclasses import dataclass
from typing import Callable

# PRE: v_pm1 \in [-1, 1]
def mapper_16bit_signed(v_pm1: float) -> int:
    return (v_pm1 + 1) / 2 * 65535 - 32768

@enum.unique
class ChannelEventKind(enum.Enum):
    NOTE_ON = 1
    NOTE_OFF = 2

@dataclass
class ChannelEvent:
    kind: ChannelEventKind
    abs_t: float
    freq: float

@dataclass
class ADSR(Callable):
    attack_ms: float
    decay_ms: float
    sustain_01: float
    release_ms: float
    hold_ms: float = 0
    sustain_ms: float = float('inf')

    def __call__(self, e: ChannelEvent, dt: float, raw: float):
        attack_s = self.attack_ms/1000
        hold_s = self.hold_ms/1000
        decay_s = self.decay_ms/1000
        if e.kind is ChannelEventKind.NOTE_ON:
            # attack stage: scale between [0, 1]
            if dt < attack_s:
                return raw*(dt/attack_s)
            # hold stage: keep at 1
            elif dt - attack_s < hold_s:
                return raw
            # decay stage: scale between [1, self.adsr.sustain_01]
            elif dt - attack_s - hold_s < decay_s:
                dt = dt - attack_s - hold_s
                a = 1 - dt / decay_s
                b = a * (1-self.sustain_01) + self.sustain_01
                return raw * b
            # sustain stage
            else:
                dt = dt - attack_s - hold_s - decay_s
                if dt < self.sustain_ms/1000:
                    return raw * self.sustain_01
                else:
                    # release stage: scale between [self.sustain_01, 0]
                    dt = dt - self.sustain_ms/1000
                    if dt < self.release_ms/1000:
                        a = 1 - dt / (self.release_ms/1000)
                        return raw * a * self.sustain_01
                    else:
                        return None
        elif e.kind is ChannelEventKind.NOTE_OFF:
            # release stage: scale between [self.adsr.sustain_01, 0]
            if dt < self.release_ms/1000:
                a = 1 - dt / (self.release_ms/1000)
                return raw * a * self.sustain_01
            else:
                return None

@enum.unique
class ChannelState(enum.Enum):
    ON = 1
    OFF = 2

def sine(freq_hz: float, time_s: float):
    return math.sin(2*math.pi*freq_hz*time_s)

def square(freq_hz: float, time_s: float):
    cycle = time_s/(1/freq_hz)
    cycle_off = cycle - math.floor(cycle)
    return -1 if cycle_off >= 0.5 else 1

def triangle(freq_hz: float, time_s: float):
    cycle = time_s*freq_hz
    cycle_off = cycle - math.floor(cycle)
    return (
        # [0, 0.25] map to [0, 1]
        cycle_off*4 if cycle_off <= 0.25
        # [0.25, 0.75] map to [1, -1]
        # -0.5: [-0.25, 0.25]
        # neg: [0.25, -0.25]
        # *4: [1, -1]
        else (cycle_off-0.5)*-4 if cycle_off <= 0.75
        # [0.75, 1] map to [-1, 0]
        # 1-: [0.25, 0]
        # neg: [-0.25, 0]
        # *4: [-1, 0]
        else (1-cycle_off)*-4
    )

def sawtooth_rising(freq_hz: float, time_s: float):
    cycle = time_s*freq_hz
    cycle_off = cycle - math.floor(cycle)
    return cycle_off*2-1
    
def sawtooth_falling(freq_hz: float, time_s: float):
    cycle = time_s*freq_hz
    cycle_off = cycle - math.floor(cycle)
    return (1-cycle_off)*2-1

@dataclass
class Channel:
    adsr: ADSR
    osc: Callable[[float, float], float] = sine
    current_sounding_freq: dict = dataclasses.field(default_factory=dict)
    current_amp_01: float = 1

    def sample(self, abs_t_s: float) -> float:
        current_sounding_freq_len = len([i for i in self.current_sounding_freq.values() if i.kind is ChannelEventKind.NOTE_ON])
        # current_sounding_freq_len = len(self.current_sounding_freq)
        if current_sounding_freq_len <= 0: return 0
        raw_note_factor = 1/current_sounding_freq_len
        raw_res = 0
        delete_freq = []
        for i in self.current_sounding_freq:
            e = self.current_sounding_freq[i]
            dt = abs_t_s - e.abs_t
            this_raw = self.osc(i, abs_t_s)

            adsr_res = self.adsr(e, dt, this_raw)
            if adsr_res is None:
                delete_freq.append(i)
            else:
                raw_res += raw_note_factor * adsr_res
                    
        for i in delete_freq:
            del self.current_sounding_freq[i]
        
        return raw_res * self.current_amp_01

    def handle_event(self, e: ChannelEvent):
        if (e.kind is ChannelEventKind.NOTE_OFF
                and self.current_sounding_freq.get(e.freq) is None):
            return
        self.current_sounding_freq[e.freq] = e

def synthesize(sample_rate_hz: int, channel_list: list[Channel], channel_event_list: list[list[ChannelEvent]]):
    res = []
    channel_list_len = len(channel_list)
    t = 0
    i = 0
    channel_event_list_copy = [i[::1] for i in channel_event_list]
    while any(channel_event_list_copy):
        s = 0
        for i in range(min(len(channel_event_list_copy), len(channel_list))):
            chnl = channel_list[i]
            evl = channel_event_list_copy[i]
            e = evl[-1]
            while evl and e.abs_t <= t:
                chnl.handle_event(e)
                evl.pop()
                if not evl: break
                e = evl[-1]
            s += (1/channel_list_len)*chnl.sample(t)
        rs = int(mapper_16bit_signed(s))
        res.append(rs&0xff)
        res.append((rs&0xff00)>>8)
        t += 1/sample_rate_hz
    
    return bytes(res)

def mk_tet(base: float, n: int):
    res = []
    for i in range(n):
        res.append(math.pow(2, i/n) * base)
    return res

a_chromatic = (
    mk_tet(110, 12)
    + mk_tet(220, 12)
    + mk_tet(440, 12)
    + mk_tet(880, 12)
    + mk_tet(880*2, 12)
)

DEFAULT_ADSR = ADSR(
    attack_ms=100,
    decay_ms=100,
    hold_ms=100,
    sustain_01=0.8,
    release_ms=300,
)
CHANNEL_LIST = [
    Channel(adsr=DEFAULT_ADSR, osc=square, current_amp_01=0.7),
]

CHANNEL_EVENT_LIST = [
    list(reversed([
        ChannelEvent(ChannelEventKind.NOTE_ON, 0, a_chromatic[0*12+2]),
        ChannelEvent(ChannelEventKind.NOTE_ON, 0, a_chromatic[1*12+2]),
        ChannelEvent(ChannelEventKind.NOTE_ON, 0, a_chromatic[1*12+5]),
        ChannelEvent(ChannelEventKind.NOTE_ON, 0, a_chromatic[1*12+12]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2, a_chromatic[0*12+2]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2, a_chromatic[1*12+2]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2, a_chromatic[1*12+5]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2, a_chromatic[1*12+12]),
        
        ChannelEvent(ChannelEventKind.NOTE_ON, 60/90*2, a_chromatic[0*12+7]),
        ChannelEvent(ChannelEventKind.NOTE_ON, 60/90*2, a_chromatic[1*12+7]),
        ChannelEvent(ChannelEventKind.NOTE_ON, 60/90*2, a_chromatic[1*12+5]),
        ChannelEvent(ChannelEventKind.NOTE_ON, 60/90*2, a_chromatic[1*12+11]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2*2, a_chromatic[0*12+7]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2*2, a_chromatic[1*12+7]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2*2, a_chromatic[1*12+5]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2*2, a_chromatic[1*12+11]),
        
        ChannelEvent(ChannelEventKind.NOTE_ON, 60/90*2*2, a_chromatic[0*12+0]),
        ChannelEvent(ChannelEventKind.NOTE_ON, 60/90*2*2, a_chromatic[1*12+0]),
        ChannelEvent(ChannelEventKind.NOTE_ON, 60/90*2*2, a_chromatic[1*12+4]),
        ChannelEvent(ChannelEventKind.NOTE_ON, 60/90*2*2, a_chromatic[1*12+11]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2*3, a_chromatic[0*12+0]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2*3, a_chromatic[1*12+0]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2*3, a_chromatic[1*12+4]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2*3, a_chromatic[1*12+11]),
        
        ChannelEvent(ChannelEventKind.NOTE_ON, 60/90*2*3, a_chromatic[0*12+5]),
        ChannelEvent(ChannelEventKind.NOTE_ON, 60/90*2*3, a_chromatic[1*12+5]),
        ChannelEvent(ChannelEventKind.NOTE_ON, 60/90*2*3, a_chromatic[1*12+4]),
        ChannelEvent(ChannelEventKind.NOTE_ON, 60/90*2*3, a_chromatic[1*12+9]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2*4, a_chromatic[0*12+5]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2*4, a_chromatic[1*12+5]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2*4, a_chromatic[1*12+4]),
        ChannelEvent(ChannelEventKind.NOTE_OFF, 60/90*2*4, a_chromatic[1*12+9]),
    ])),
]

output_file = wave.open('test.wav', 'wb')
output_file.setsampwidth(2)
output_file.setnchannels(1)
output_file.setframerate(44100)
s = synthesize(44100, CHANNEL_LIST, CHANNEL_EVENT_LIST)
output_file.writeframes(s)
output_file.close()