import librosa
import pyloudnorm as pyln
import numpy as np
import math
import random

from textual.containers import Vertical
from textual.widgets import Sparkline
#from textual_plot import HiResMode, PlotWidget

from ui.widgets.chunk_types.chunk import Chunk



class AudioChunk(Chunk):
    DEFAULT_CSS = """
    AudioChunk {
        border: tab $primary;
        padding: 1 0;
        Sparkline {
            height: 1fr;
        }
    }
    """

    def __init__(self, audio_data: np.ndarray, sample_rate: int, chunk_name: str = "Sample", bar_pos: float = 0.0):
        super().__init__(chunk_name, bar_pos)
        
        self.audio = audio_data
        self.sample_rate = sample_rate
        
        self.num_channels = None
        if len(self.audio.shape) == 1:
            self.num_samples = self.audio.shape[0]
            self.num_channels = 1
        else:
            self.num_samples = self.audio.shape[1]
            self.num_channels = self.audio.shape[0]
        
        self.meter = pyln.Meter(self.sample_rate)
        self.loudness_values = []
        
        self.calculate_size()
        
    def calculate_size(self):
        self.styles.width = round((self.num_samples / self.sample_rate) / self.app.zoom_level)
    
    def on_mount(self):
        for plot in self.query(Sparkline):
            plot: Sparkline = plot
            #plot: PlotWidget = plot # just for type checking

            #plot.margin_top = 0
            #plot.margin_left = 0
            #plot.margin_bottom = 0
            
            num_values = len(list(range(0, int(self.num_samples), int(self.sample_rate * 0.05))))
            #x = range(num_values)

            y = []
            if self.num_channels == 1:
                # get rms
                rms = librosa.feature.rms(y=self.audio)[0]
                
                # split into N chunks
                rms = np.array([
                    np.mean(chunk) if len(chunk) > 0 else 0.0
                    for chunk in np.array_split(rms, num_values) 
                ])
                
                # get the decibel values from that
                y = list(librosa.amplitude_to_db(rms, ref=np.min))
                
            else:
                # get rms
                rms = librosa.feature.rms(y=self.audio[int(plot.id[-1])])[0]
                
                # split into N chunks
                rms = np.array([
                    np.mean(chunk) if len(chunk) > 0 else 0.0
                    for chunk in np.array_split(rms, num_values)
                ])
                
                # get the decibel values from that
                y = librosa.amplitude_to_db(rms, ref=np.min)

            plot.data = list(y)

            """print(x)
            plot.bar(
                    x,
                    y,
                    1.0,
                    bar_style=self.app.theme_variables["primary"],
                    hires_mode=HiResMode.BRAILLE
                )
                
                
            plot.set_xticks([])
            plot.set_yticks([])
            plot.set_ylimits(ymin=0)"""
        
    def compose(self) -> ComposeResult:
        
        
        # account for multi-track audio and mono audio
        if self.num_channels > 1:
            # loop over each channel in the audio and display it seperately
            for channel in range(self.num_channels):
                
                """samples = []
                
                for sample in range(0, self.num_samples, int(self.sample_rate*0.1)):
                    samples.append(self.audio[channel, sample])"""
                
                yield Sparkline(id=f"channel-{channel}")
                #yield PlotWidget(allow_pan_and_zoom=False, id=f"channel-{channel}")
                
                
        else:
            # just display the one channel
            """samples = []
            
            for sample in range(0, self.num_samples, int(self.sample_rate*0.1)):
                samples.append(self.audio[sample])"""
                
            #yield PlotWidget(allow_pan_and_zoom=False)
            yield Sparkline()