from PyQt5.QtGui import QIcon, QPixmap
from PyQt5.QtWidgets import (
    QWidget, QLabel, QHBoxLayout, QVBoxLayout, QApplication,
    QMainWindow, QGridLayout, QDialog, QFrame, QTabWidget,
    QScrollArea
)

import numpy as np

import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolBar
from matplotlib.colors import LogNorm, BoundaryNorm

import settings as stg


class PlotNoiseWindow(QDialog):
    def __init__(self, parent=None):
        super(PlotNoiseWindow, self).__init__(parent)

        self.setGeometry(400, 200, 700, 500)

        self.setWindowTitle("Noise field")

        self.verticalLayout_Main = QVBoxLayout()
        self.setLayout(self.verticalLayout_Main)

        self.tab = QTabWidget()
        self.verticalLayout_Main.addWidget(self.tab)

        self.tabs = []

        for i in range(len(stg.filename_BS_raw_data)):
            self.tabs.append(QWidget())

            self.tab.addTab(self.tabs[i], stg.filename_BS_raw_data[i])

            verticalLayout = QVBoxLayout()
            self.tabs[i].setLayout(verticalLayout)

            fig, ax = plt.subplots(
                nrows=stg.freq[i].shape[0], ncols=1, layout='constrained'
            )

            if stg.BS_noise_raw_data[i].shape != (0,):
                for freq_ind, freq_val in enumerate(stg.freq[i]):
                    ax[freq_ind].cla()

                    val_min = np.nanmin(stg.BS_noise_raw_data[i][freq_ind, :, :])
                    val_max = np.nanmax(stg.BS_noise_raw_data[i][freq_ind, :, :])

                    if val_min != val_max:
                        val_min = 0
                        val_max = 1e-5

                    pcm = ax[freq_ind].pcolormesh(
                        stg.time_noise[i][freq_ind, :],
                        -stg.depth_noise[i][freq_ind, :],
                        stg.BS_noise_raw_data[i][freq_ind, :, :],
                        cmap='hsv'
                    )

                    ax[freq_ind].text(
                        1, .70, stg.freq_text[i][freq_ind],
                        fontsize=10, fontweight='bold',
                        fontname='DejaVu Sans', c='black', alpha=0.5,
                        horizontalalignment='right', verticalalignment='bottom',
                        transform=ax[freq_ind].transAxes
                    )

                fig.supxlabel('Time (sec)', fontsize=10)
                fig.supylabel('Depth (m)', fontsize=10)

                cbar = fig.colorbar(
                    pcm, ax=ax[:], shrink=1, location='right'
                )
                cbar.set_label(label='Noise signal (V)', rotation=270, labelpad=8)
                fig.canvas.draw_idle()

            canvas = FigureCanvas(fig)
            toolbar = NavigationToolBar(canvas, self)

            scroll = QScrollArea()
            scroll.setWidget(canvas)

            verticalLayout.addWidget(toolbar)
            verticalLayout.addWidget(scroll)