"""made by Mohammad Erfan Karami
github: https://github.com/erfan-ops

this package is used to create play and save sound files
it has some basic sound waves although you can add your own and you can modify the package.

it stores the sound waves as numpy arrays and uses pyaudio for playback
uses scipy.io.wavfile to read eav files because the wave module can read the float32 bit wav files
and finally uses matplotlib to visualize the waves"""


import pyaudio
import numpy as np
from typing import Callable, Iterable, Literal
import scipy.io.wavfile as wf
import matplotlib.pyplot as plt


# types #
SoundBuffer = np.ndarray[np.any, np.dtype[np.float32|np.int16|np.uint8]]
Wave = Callable[[float, float, float], np.ndarray[np.any, np.dtype]]


# its just a dictionary
class Notes(dict[str, float]):
    def __init__(self):
        super().__init__()
    
    def __getitem__(self, __key: str) -> float:
        return super().__getitem__(__key.upper())
    
    def __setitem__(self, __key: str, __value: float) -> None:
        return super().__setitem__(__key.upper(), __value)


#-- sound waves --#
class Sounds:
    def __init__(self, sample_rate: int=48000, dtype:np.float32|np.int16|np.uint8=np.float32) -> None:
        self.default_sample_rate: int = sample_rate
        self.existed_notes: dict[str, tuple] = {}
        self.total_cached: int = 0
        
        self.tau = 2*np.pi
        self.rev_pi = 1/np.pi
        self.two_oper_pi = 2/np.pi
        self.four_over_pi = 4/np.pi
        self.eight_over_pi_sqr = 8/np.pi**2
        
        self.dtype = dtype
        
    # caches the waves so if you had to generate a note multiple times its not gonna take long to execute
    # you can't use the lru_cache from the functools module because waves are stored as numpy arrays and numpy arrays are not hashable
    # so this decorator will save a tuple of that numpy array in self.existed_notes and convert it back to a numpy array each time you want to use that wave
    def cache_wave(wave_type:str) -> SoundBuffer:
        """returns the cached value if available, otherwise caches the returned wave"""
        def decorator(func):
            def wrapper(self, freq:float, dur:float, vol: float):
                name = f"{wave_type}|{freq}|{dur}|{vol}"
                if name in self.existed_notes.keys():
                    self.total_cached += 1
                    return np.array(self.existed_notes[name], dtype=self.dtype)
                
                result = func(self, freq, dur, vol)
                self.existed_notes[name] = tuple(result)
                
                return result
            return wrapper
        return decorator

    @cache_wave("sine")
    def sine_wave(self, freq:float, dur:float, vol: float) -> SoundBuffer:
        """creates a sine wave based on the given frequency, duration and amplituse or volume
        returns a numpy array"""
        
        wave = (vol * np.sin(self.tau * np.arange(self.default_sample_rate * dur) * freq / self.default_sample_rate)).astype(self.dtype)
        return wave
    
    @cache_wave("sqr")
    def square_wave(self, freq:float, dur:float, vol: float) -> SoundBuffer:
        """creates a square wave based on the given frequency, duration and amplitude or volume\n
        a square wave is typically generated by adding all the odd harmonics (3, 5, 7, 9...) to a fundamental frequency\n
        each harmonic is added with a little bit less volume which causes the amplitude to change a bit so at the end the amplitude is multiplied by 4/pi to fix it
        returns a numpy array"""
        
        # -- creating the fundumental note --#
        buf = np.sin(self.tau * np.arange(self.default_sample_rate * dur) * freq / self.default_sample_rate)
        for h in range(3, 1000, 2):
            f = freq*h
            #-- break if the frequency is greater than 20,000 because human ear can' hear it and the arent many headphones that can play more than this anyways --#
            if f > 20_000:
                break
            #-- adding the harmonics --#
            buf += np.sin(self.tau * np.arange(self.default_sample_rate * dur) * f / self.default_sample_rate)/h
        
        wave = ((self.four_over_pi)*vol*buf).astype(self.dtype)
        return wave
    
    @cache_wave("tri")
    def triangle_wave(self, freq:float, dur:float, vol: float) -> SoundBuffer:
        """creates a triangle wave based on the given frequency, duration and amplitude or volume\n
        a triangle wave is typically generated by adding every second odd harmonic (5, 9, 13...) to a fundamental frequency\n
        each harmonic is added with a less volume which causes the amplitude to change a bit so at the end the amplitude is multiplied by '8/pi^2' to fix it\n
        returns a numpy array"""
        
        buf = np.sin(self.tau * np.arange(self.default_sample_rate * dur) * freq / self.default_sample_rate)
        for h in range(1, 500):
            harmonic = (2*h+1)
            f = freq*harmonic
            if f > 20000:
                break
            buf += (-1)**h * harmonic**(-2) * np.sin(self.tau * np.arange(self.default_sample_rate * dur) * f / self.default_sample_rate)
        
        wave = ((self.eight_over_pi_sqr)*vol*buf).astype(self.dtype)
        return wave
    
    @cache_wave("fast_tri")
    def fast_triangle_wave(self, freq:float, dur:float, vol: float) -> SoundBuffer:
        return (self.two_oper_pi * vol * np.arcsin(np.sin(self.tau * np.arange(self.default_sample_rate * dur) * freq / self.default_sample_rate))).astype(self.dtype)
    
    @cache_wave("saw")
    def sawtooth_wave(self, freq:float, dur:float, vol: float) -> SoundBuffer:
        """returns a numpy array,
        creates a sawtooth wave based on the given frequency, duration and amplitude or volume\n
        a sawtooth wave is typically generated by adding every harmonic to a fundamental frequency with the odd harmonics being negative (2, -3, 4, -5...)\n
        each harmonic is added with less volume which causes the amplitude to change a bit so at the end we do this
        >>> (vol * (0.5 - self.rev_pi*buf))"""
        
        buf = (-1) * np.sin(self.tau * np.arange(self.default_sample_rate * dur) * freq / self.default_sample_rate)
        for h in range(2, 1000):
            f = freq*h
            if f > 20000:
                break
            buf += (-1)**h * (np.sin(self.tau * f * np.arange(self.default_sample_rate * dur) / self.default_sample_rate)/h)
        
        wave = (vol * ((-self.two_oper_pi)*buf)).astype(np.float32)
        return wave
    
    @cache_wave("revsaw")
    def reversed_sawtooth_wave(self, freq:float, dur:float, vol: float) -> SoundBuffer:
        """creates a reversed sawtooth wave based on the given frequency, duration and amplitude or volume\n
        a reversed sawtooth wave is typically generated by adding every harmonic to a fundamental frequency with the odd harmonics being negative (2, -3, 4, -5...)\n
        each harmonic is added with less volume which causes the amplitude to change a bit so at the end we do this
        >>> ((2*vol/np.pi) * (0.5 - self.rev_pi*buf))"""
        
        buf = (-1) * np.sin(self.tau * np.arange(self.default_sample_rate * dur) * freq / self.default_sample_rate)
        for h in range(2, 1000):
            f = freq*h
            if f > 20000:
                break
            buf += (-1)**h * (np.sin(self.tau * f * np.arange(self.default_sample_rate * dur) / self.default_sample_rate)/h)
        
        wave = (vol * self.two_oper_pi * buf).astype(np.float32)
        return wave
    
    @cache_wave("organ")
    def organ(self, freq:float, dur:float, vol: float) -> SoundBuffer:
        """creates an organ like sound based on the given frequency, duration and amplitude or volume"""
        
        pi_2_samples_num = self.tau * np.arange(self.default_sample_rate * dur)
        buf = vol * np.sin(pi_2_samples_num * freq / self.default_sample_rate)
        buf += vol*0.9 * np.sin(pi_2_samples_num * freq*2 / self.default_sample_rate)
        buf += vol * np.sin(pi_2_samples_num * freq*4 / self.default_sample_rate)
        buf += vol*0.6 * np.sin(pi_2_samples_num * freq*6 / self.default_sample_rate)
        buf += vol*0.7 * np.sin(pi_2_samples_num * freq*16 / self.default_sample_rate)
        buf += vol*0.5 * np.sin(pi_2_samples_num * freq*20 / self.default_sample_rate)
        buf += vol*0.3 * np.sin(pi_2_samples_num * freq*24 / self.default_sample_rate)

        wave = (buf/7).astype(self.dtype)
        return wave
    
    @cache_wave("marimb")
    def marimba(self, freq:float, dur:float, vol: float) -> SoundBuffer:
        """creates a "not even close to marimba" sound based on the given frequency, duration and amplitude or volume\n
        warning!: very annoying sound"""
        
        pi_2_samples_num = self.tau * np.arange(self.default_sample_rate * dur)
        buf = vol * np.sin(pi_2_samples_num * freq / self.default_sample_rate)
        buf += vol*0.75 * np.sin(pi_2_samples_num * freq*10 / self.default_sample_rate)
        buf += vol/2 * np.sin(pi_2_samples_num * freq*20 / self.default_sample_rate)
        buf += vol/4 * np.sin(pi_2_samples_num * freq*30 / self.default_sample_rate)
        
        wave = (buf/4).astype(self.dtype)
        return wave
    
    @cache_wave("tst")
    def test(self, freq:float, dur:float, vol: float) -> SoundBuffer:
        buf = np.sin(self.tau * np.arange(self.default_sample_rate * dur) * freq / self.default_sample_rate)
        for h in range(1, 1000):
            f = freq*h
            if f > 20000:
                break
            buf += (-1)**h * (np.sin(self.tau * np.arange(self.default_sample_rate * dur) * f / self.default_sample_rate)/h)
        
        wave = (vol * (0.5 - self.rev_pi*buf)).astype(self.dtype)
        return wave
    
    
    def array_to_tuple(self, np_array):
        """Iterates recursivelly."""
        try:
            return tuple(self.array_to_tuple(_) for _ in np_array)
        except TypeError:
            return np_array


class Export:
    def __init__(self) -> None:
        pass
    
    # converts a float32 array to an int type
    def float2pcm(self, sig: np.ndarray, dtype: Literal["int32", "int16", "uint8"]="int16") -> SoundBuffer:
        sig = np.asarray(sig)
        if sig.dtype.kind != 'f':
            raise TypeError("'sig' must be a float array")
        dtype = np.dtype(dtype)
        if dtype.kind not in 'iu':
            raise TypeError("'dtype' must be an integer type")

        i = np.iinfo(dtype)
        abs_max: int = 2 ** (i.bits - 1)
        offset = i.min + abs_max
        return (sig * abs_max + offset).clip(i.min, i.max).astype(dtype)
    
    # converts an int type array to a float32
    def pcm2float(self, sig: np.ndarray, dtype: Literal["float32"]="float32") -> np.ndarray[np.any, np.floating]:
        sig = np.asarray(sig)
        if sig.dtype.kind not in 'iu':
            raise TypeError("'sig' must be an array of integers")
        dtype = np.dtype(dtype)
        if dtype.kind != 'f':
            raise TypeError("'dtype' must be a floating point type")

        i = np.iinfo(sig.dtype)
        abs_max: int = 2 ** (i.bits - 1)
        offset = i.min + abs_max
        return (sig.astype(dtype) - offset) / abs_max
        
    # exports to .erfan file which is completely useless but you can play them with the app.py file adn you can convert them to wav file which is actualy useful
    def export_to_erfan(self, file_name:str, buffer: np.ndarray|bytes, sample_rate: int, dtype, channels: int) -> None:
        if type(buffer) != bytes:
            dtype = buffer.dtype
            if dtype == np.uint8 or dtype == "uint8":
                audio_format = 1
            elif dtype == np.int16 or dtype == "int16":
                audio_format = 2
            elif dtype == "int24":
                audio_format = 3
            elif dtype == np.float32 or dtype == "float32":
                audio_format = 4
            else:
                raise ValueError(f"Incorrect dtype: {dtype}\nvalid options are: \"uint8, int16, int24, float32\"")
                
        
        if not sample_rate:
            sample_rate = self.default_sample_rate
        
        with open(f"{file_name}.erfan", "wb") as f:
            f.write(sample_rate.to_bytes(4, "little"))
            f.write(audio_format.to_bytes(2, "little"))
            f.write(channels.to_bytes(2, "little"))
            
            if type(buffer) == np.ndarray:
                buffer = buffer.tobytes()
            
            f.write(buffer)
    
    
    def export_to_wav(self, file_path:str, buffer: np.ndarray|bytes, sample_rate: int, dtype) -> None:
        if type(buffer) == bytes:
            buffer = np.frombuffer(buffer[8:], dtype=dtype)
        
        wf.write(f"{file_path}.wav", sample_rate, buffer)
    
    # reads and return the bytes containing the sound from an erfan file
    def read_from_erfan(self, file_name: str) -> bytes:
        with open(file_name, "rb") as f:
            data = f.read()
            sample_rate = int.from_bytes(data[0:4], "little")
            sampwidth = int.from_bytes(data[4:6], "little")
            channels = int.from_bytes(data[6:8], "little")
            
            self.stream = self.AUDIO_OBJECT.open(format=self.AUDIO_OBJECT.get_format_from_width(sampwidth),
                                                 channels=channels,
                                                 rate=sample_rate,
                                                 output=True)
            return data[8:]
    
    
    def erfan_to_wav(self, file_name: str, dir: str="", dtype=np.float32) -> None:
        file_short_name = file_name.removesuffix(".erfan")
        with open(file_name, "rb") as f:
            data = f.read()
            sample_rate = int.from_bytes(data[0:4], "little")
            data = np.frombuffer(data[8:], dtype=dtype)
        
        wf.write(f"{dir}{file_short_name}.wav", sample_rate, data)
    
    
    def wav_to_erfan(self, file_name: str) -> None:
        file_short_name = file_name.removesuffix(".wav")
        with open(file_name, "rb") as f:
            data = f.read()
            sample_rate = int.from_bytes(data[24:28], "little")
            data = data[42:]
        
        with open(f"{file_short_name}.erfan", "wb") as f:
            f.write(sample_rate.to_bytes(4, "little"))
            f.write(sample_rate.to_bytes(2, "little"))
            f.write(sample_rate.to_bytes(2, "little"))
            f.write(data)

    # converts a float32 wav file to an int type
    def wav_float32_to_int(self, file_path: str, dtype: Literal["int32", "int16", "uint8"]="int16"):
        with open(file_path, "rb") as f:
            sample_rate, data = wf.read(file_path)
            if data.dtype != "float32":
                raise f"This function is only used to conver \"float32\" dtype to \"int\" not <{data.dtype}>"
            
            wf.write(f"{file_path}.wav", sample_rate, self.float2pcm(data, dtype))

    
    def get_sampwidth_from_str(self, dtype: Literal["float32", "int24", "int16", "uint8"]) -> int:
        match dtype:
            case "uint8":
                width = 1
            case "int16":
                width = 2
            case "int24":
                width = 3
            case "float32":
                width = 4
            case _:
                return 0
        
        return width


class Music(Sounds, Export):
    def __init__(self, middle_a=440, tempo=60, sample_rate=48000, tune=12) -> None:
        super().__init__(sample_rate)
        self.sample_rate: int = sample_rate
        self.AUDIO_OBJECT = pyaudio.PyAudio() # making an audio object that we,ll need to play sounds
        self.TUNE = tune                       # using 12 tone equal temperament
        self.MIDDLE_A = middle_a                  # middle a has a frecuency of 440 hertz
        self._NOTES: Notes = Notes()                 # a directory that has all the notes that we can use in it and R is just for rest
        self.NOTE_NAMES: list[str] = list("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
        self.FULL_NOTE_NAMES = []
        self.DEFAULT_DURATION = 1
        self.DEFAULT_VOLUME = 0.5
        self.tempo = tempo
        self.channels: int = 1
    
    
    # you have to use this function to initialize before playing any note.
    def init(self) -> None:
        super().__init__(self.sample_rate)
        self.SEMI_TONE: float = np.power(2, (1/self.TUNE)) # a semitone in a 12 equal temperament is 2 to the power of 1/12, that means in a 6 equal temperament system a semitone is 2^(1/6)
        self.generate_note_names()
        self.assign_middle_a()
        self.assign_frequencies()
        
        self.stream = self.config_stream(samp_width="float32",
                                         channels=self.channels,
                                         sample_rate=self.sample_rate)
    
    
    def config_stream(self, samp_width: int, channels: int, sample_rate: int) -> pyaudio.Stream:
        if type(samp_width) == str:
            samp_width = self.get_sampwidth_from_str(dtype=samp_width)
        
        return self.AUDIO_OBJECT.open(format=self.AUDIO_OBJECT.get_format_from_width(samp_width),
                                      channels=channels,
                                      rate=sample_rate,
                                      output=True)
    
    
    def generate_note_names(self) -> None:
        # the number of note names has to be the same as the tune so if its less its not going to work.
        if self.FULL_NOTE_NAMES and (len(self.FULL_NOTE_NAMES) < self.TUNE):
            raise Exception(f"Not enough note names for base {self.TUNE} scale")

        for n in self.NOTE_NAMES:
            if len(self.FULL_NOTE_NAMES) == self.TUNE:
                break
            self.FULL_NOTE_NAMES.append(n)
            
            # we don't need B# and E# because we have enharmonic notes for them,
            # but if you want to use those you can by changing the notes manually.
            if n != "B" and n != "E":
                self.FULL_NOTE_NAMES.append(f"{n}#")
    
    
    # if you want to change the middle a in the middle of a song you need to use this function
    def assign_middle_a(self, new_frequency:int|float=None) -> None:
        if new_frequency:
            self.MIDDLE_A = new_frequency
        self.BASS_A = self.MIDDLE_A
        self._hbl = 1
        
        while self.BASS_A >= 20:
            self.BASS_A /= 2
            self._hbl += 1
    
    # you need to call the "assign_frequencies" function after it to change all the note's frecuencies based on the new middle a
    # although you can just change the self.middle_a value and use init after it, which will also call this function.
    
    
    # this function will create a dictionary of all the notes with a range of 20- to 20'000 frequencies.
    # it is depended on the middle_a and the tune that you're using.
    def assign_frequencies(self) -> None:
        octave = 0
        frequency = 0
        
        # this loop will create the directory and the name of notes are formatted like this "'{note}{octave}':frecuency"
        while frequency <= 20000:
            for i in range(self.TUNE):
                frequency = self.BASS_A * np.power(self.SEMI_TONE, i) * (np.power(2, octave)).astype(np.float32)
                
                if frequency > 20000:
                    break
                
                elif frequency >= 20:
                    self._NOTES[self.FULL_NOTE_NAMES[i%self.TUNE]+str(octave-(self._hbl-5))] = frequency

            octave += 1
    
    
    def fix_duration(self, duration:str|float=0) -> float:
        if not duration:
            duration = self.DEFAULT_DURATION
        
        elif str(duration)[-1].lower() != "s":
            duration = 60 / self.tempo * float(duration)
        else:
            duration = float(duration[:-1])
        
        return duration
    
    
    def fix_volume(self, volume:float=0) -> float:
        if volume == 0:
            volume = self.DEFAULT_VOLUME
        volume = abs(volume)

        if volume >= 1:
            volume = 1
        
        return volume
    
    
    # creates a note buffer but doesn't turn it into bytes
    def generate_note_buffer(self, note:str|float|int, wave_type: Wave, duration:str|float=0, volume:float=0):
        duration = self.fix_duration(duration)
        volume = self.fix_volume(volume)
        try:
            note = float(note)
            f = note
        except ValueError:
            note = note.upper()
            f = self._NOTES[note]
        
        buf = wave_type(f, duration, volume)
        return buf
    
    
    # creates a chord buffer but doesn't turn it into bytes
    def generate_chord_buffer(self, notes: Iterable, wave_type: Wave, duration:str|float=0, volume:float=0):
        buf = self.generate_note_buffer(notes[0], wave_type, duration, volume)
        for n in notes[1:]:
            buf += self.generate_note_buffer(n, wave_type, duration, volume)
        return (buf/len(notes))
    
    
    # creates a note buffer and plays it
    def play(self, note:str|float|int, wave_type: Wave, duration:str|float=0, volume:float=0) -> None:
        buf: bytes = self.generate_note_buffer(note, wave_type, duration, volume).tobytes()
    
        self.play_buffer(buf)
    
    
    # plays a buffer
    def play_buffer(self, buffer: np.ndarray|bytes) -> None:
        if type(buffer) != bytes:
            buffer = buffer.tobytes()
        
        self.stream.write(buffer)
    
    
    # created a chord sound buffer and plays it
    def play_chord(self, notes: Iterable, wave_type: Wave, duration:str|float=0, volume:float=0) -> None:
        buf = self.generate_chord_buffer(notes, wave_type, duration, volume).tobytes()
        
        self.play_buffer(buf)
    
    
    def play_erfan(self, file_name: str) -> None:
        self.play_buffer(self.read_from_erfan(file_name))
    
    
    def play_wav(self, file_name: str) -> None:
        sample_rate, data = wf.read(file_name)
        match data.dtype:
            case "uint8":
                sampwidth = 1
            case "int16":
                sampwidth = 2
            case "int24":
                sampwidth = 3
            case "float32":
                sampwidth = 4
            
            #-- self.AUDIO_OBJECT is a pyaudio.PyAudio                        --#
            #-- and get_format_from_width(2) will just return pyaudio.PaInt16 --#
            #-- (4 will return a PaFloat32)                                   --#
            
        self.stream = self.AUDIO_OBJECT.open(format=self.AUDIO_OBJECT.get_format_from_width(sampwidth),
                                             channels=data.ndim,
                                             rate=sample_rate,
                                             output=True)
        self.play_buffer(data)
    
    
    def visualize_sound(self, wave: SoundBuffer, sample_rate: int) -> None:
        times = np.linspace(0, wave.size/sample_rate, wave.size)
        duration = wave.size / sample_rate
        
        plt.figure(figsize=(15, 5))
        plt.plot(times, wave)
        plt.title('wave:')
        plt.ylabel('Signal Value')
        plt.xlabel('Time (s)')
        plt.xlim(0, duration)
        plt.show()
    
    
    def show_frequency_spectrum(self, wave: SoundBuffer, sample_rate: int) -> None:
        duration = wave.size / sample_rate
        
        plt.figure(figsize=(15, 5))
        plt.specgram(wave, Fs=sample_rate, vmin=-20, vmax=50)
        plt.title('wave')
        plt.ylabel('Frequency (Hz)')
        plt.xlabel('Time (s)')
        plt.xlim(0, duration)
        plt.colorbar()
        plt.show()
    
    
    # this will close the stream and terminate the audio_object (basicly closes everything)
    def done(self) -> None:
        self.stream.stop_stream()
        self.stream.close()
        self.AUDIO_OBJECT.terminate()