# dr14meter: compute the DR14 value of the given audio files
# Copyright (C) 2024  pe7ro
#
# dr14_t.meter: compute the DR14 value of the given audiofiles
# Copyright (C) 2012  Simone Riva
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import numpy as np
from dr14meter.audio_math import *
from dr14meter.plot.my_time_formatter import *
from dr14meter.out_messages import *

try:
    import matplotlib
    import matplotlib.pyplot as pyplot
    import matplotlib.mlab as mlab
    from matplotlib.widgets import *
    from pylab import *
except:
    pass


def on_select(vmin, vmax):
    dr14_log_debug("on_select: %f , %f " % (vmin, vmax))
    if vmin == vmax:
        return
    plot_track_classic(plot_str=on_select.plot_str,
                       start_time=vmin, end_time=vmax)
    on_select.plot_str.plot()


def wheel_moved(event):
    dr14_log_debug("mouse_pressed: step: %d " % event.step)

    (start_time, end_time) = wheel_moved.plot_str.move(event.step)
    plot_track_classic(plot_str=wheel_moved.plot_str,
                       start_time=start_time, end_time=end_time)
    wheel_moved.plot_str.plot()


def mouse_pressed(event):

    dr14_log_debug("mouse_pressed: ")

    if event.button == 3:
        dr14_log_debug("mouse_pressed: button 3")

        (start_time, end_time) = mouse_pressed.plot_str.zoom_out()
        plot_track_classic(plot_str=mouse_pressed.plot_str,
                           start_time=start_time, end_time=end_time)
        mouse_pressed.plot_str.plot()

        dr14_log_debug("mouse_pressed: start - end:  %f %f " %
                       (start_time, end_time))


class PltTrackStruct:

    def __init__(self, plot_mode='fill', Y=[], Fs=44100, sz=0, ch=0):
        self.ch = ch
        self.tot_time = 0

        self.t = None
        self.tb = None
        self.start_time = 0
        self.end_time = 0

        self.start_block = 0
        self.end_block = 0

        self.mp = None
        self.mn = None

        self.Y = Y
        self.Fs = Fs
        self.plot_mode = plot_mode  # "fill" or "curve"
        self.ax = None
        self.rebuild = False
        self.first_sample = 0
        self.sz_section = 0
        self.lines = []
        self.fig = None

    def start(self):
        self.plot()
        pyplot.show()

    def zoom_out(self):
        start_time = self.start_time
        end_time = self.end_time

        m = start_time + (end_time - start_time) / 2.0
        new_tot_time = 2.0 * (end_time - start_time)

        start_time = m - new_tot_time / 2.0
        end_time = m + new_tot_time / 2.0

        return (start_time, end_time)

    def move(self, sign):

        factor = 0.05

        delta = self.end_time - self.start_time

        start_time = self.start_time + sign * delta * factor
        end_time = self.end_time + sign * delta * factor

        if start_time < 0 or end_time > self.tot_time:
            start_time = self.start_time
            end_time = self.end_time

        return (start_time, end_time)

    def plot(self):

        new_flag = False

        if self.fig == None:
            self.fig = pyplot.figure()
            self.ax = [i for i in range(self.ch)]
            self.lines = [i for i in range(self.ch)]
            self.span = []
            new_flag = True

        for j in range(self.ch):

            if new_flag:
                self.ax[j] = self.fig.add_subplot(210 + j + 1)

            if type(self.lines[j]) is matplotlib.collections.PolyCollection:
                self.lines[j].remove()
            else:
                while (not new_flag) and len(self.lines[j]) > 0:
                    l = self.lines[j].pop(0)
                    l.remove()

            if self.plot_mode == 'fill':
                self.lines[j] = self.ax[j].fill_between(self.tb[self.start_block:self.end_block],
                                                        self.mp[
                                                            self.start_block:self.end_block, j],
                                                        self.mn[self.start_block:self.end_block, j])
            else:
                self.lines[j] = self.ax[j].plot(
                    self.t, self.Y[self.first_sample: self.first_sample + self.sz_section, j], 'b')

            self.ax[j].axis([self.start_time, self.end_time, -1.0, 1.0])

            delta_time = self.end_time - self.start_time

            if delta_time < 1.5:
                milli_sec = True
            else:
                milli_sec = False

            self.ax[j].xaxis.set_major_formatter(
                MyTimeFormatter(milli_sec=milli_sec))

            self.ax[j].grid(True)

            pyplot.title("Channel %d" % (j + 1))
            pyplot.xlabel('Time [min:sec]')
            pyplot.ylabel('Amplitude')

            if new_flag:
                onsel = on_select
                onsel.plot_str = self
                self.span.append(SpanSelector(
                    self.ax[j], onsel, 'horizontal', button=1))

                m_p = mouse_pressed
                m_p.plot_str = self
                connect('button_press_event', m_p)

                w_m = wheel_moved
                w_m.plot_str = self
                connect('scroll_event', w_m)

        self.fig.canvas.draw()


def plot_track_classic(Y=None, Fs=None, plot_str=None, utime=0.02, time_lim=5, start_time=0.0, end_time=-1.0):

    if Y is None and Fs is None and plot_str is None:
        raise

    if plot_str is None and Y is not None and Fs is not None:
        s = Y.shape

        if len(Y.shape) > 1:
            ch = s[1]
        else:
            ch = 1

        plot_str = PltTrackStruct(ch=ch)
        plot_str.Y = Y
        plot_str.Fs = Fs
        plot_str.tot_time = s[0] * 1.0 / Fs

    elif plot_str != None:
        Fs = plot_str.Fs
        Y = plot_str.Y

        s = Y.shape
        if len(Y.shape) > 1:
            ch = s[1]
        else:
            ch = 1
    else:
        raise

    sz_section = s[0]
    first_sample = 0

    if end_time > s[0] * 1.0 / Fs:
        end_time = s[0] * 1.0 / Fs

    if start_time < 0:
        start_time = 0.0

    if start_time >= 0.0 and end_time > start_time:
        sz_section = int((end_time - start_time) * Fs)
        first_sample = int(start_time * Fs)
        plot_str.start_time = start_time
        plot_str.end_time = end_time
    else:
        plot_str.start_time = 0.0
        plot_str.end_time = sz_section * 1.0 / Fs

    block_len = utime

    samples_block = int(block_len * Fs)
    sz = int(sz_section / samples_block)

    curr_sample = first_sample

    if plot_str.mp is not None and plot_str.mn is not None and sz_section >= time_lim * Fs:
        plot_str.plot_mode = "fill"
        plot_str.start_block = int(plot_str.start_time / block_len)
        plot_str.end_block = int(plot_str.end_time / block_len)

        return plot_str

    if sz_section > time_lim * Fs:
        plot_str.plot_mode = "fill"
        plot_str.mp = zeros((sz, ch))
        plot_str.mn = zeros((sz, ch))
        plot_str.tb = start_time + np.arange(0, sz) * block_len

        for i in range(sz):
            plot_str.mp[i, :] = np.max(
                Y[curr_sample:curr_sample + samples_block, :], 0)
            plot_str.mn[i, :] = np.min(
                Y[curr_sample:curr_sample + samples_block, :], 0)
            curr_sample = curr_sample + samples_block

        plot_str.start_block = 0
        plot_str.end_block = sz

    else:
        plot_str.plot_mode = "curve"
        plot_str.t = start_time + np.arange(0, sz_section) * 1.0 / Fs
        plot_str.first_sample = first_sample
        plot_str.sz_section = sz_section

    return plot_str
