from Model.AquascatDataLoader import RawAquascatData

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm


class AcousticDataLoader:

    def __init__(self, path_BS_raw_data: str):

        self.path_BS_raw_data = path_BS_raw_data
        # --- Load Backscatter acoustic raw data with RawAquascatData class ---

        self._data_BS = RawAquascatData(self.path_BS_raw_data)
        self._BS_raw_data = np.swapaxes(self._data_BS.V, 0, 1)

        self._freq = self._data_BS.Freq
        self._freq_text = self._data_BS.freqText

        self._r = np.repeat(np.transpose(self._data_BS.r), self._freq.shape[0], axis=0)

        self._time = np.repeat(
            np.transpose(np.array([t / self._data_BS.PingRate for t in range(self._data_BS.NumProfiles)])[:, np.newaxis]),
            self._freq.shape[0], axis=0)

        self._date = self._data_BS.date.date()
        self._hour = self._data_BS.date.time()
        self._nb_profiles = [self._data_BS.NumProfiles]*self._freq.shape[0]
        self._nb_profiles_per_sec = [self._data_BS.ProfileRate]*self._freq.shape[0]
        self._nb_cells = [self._data_BS.NumCells]*self._freq.shape[0]
        self._cell_size = [self._data_BS.cellSize]*self._freq.shape[0]
        self._pulse_length = [self._data_BS.TxPulseLength]*self._freq.shape[0]
        self._nb_pings_per_sec = [self._data_BS.PingRate]*self._freq.shape[0]
        self._nb_pings_averaged_per_profile = [self._data_BS.Average]*self._freq.shape[0]
        self._kt = self._data_BS.Kt.tolist()
        self._gain_rx = self._data_BS.RxGain.tolist()
        self._gain_tx = self._data_BS.TxGain.tolist()

    def reshape_BS_raw_data(self):
        BS_raw_cross_section = np.reshape(self._BS_raw_data,
                                          (self._r.shape[1] * self._time.shape[1], self._freq.shape[0]),
                                          order="F")
        return BS_raw_cross_section

    def reshape_r(self):
        r = np.zeros((self._r.shape[1] * self._time.shape[1], self._freq.shape[0]))
        for i, _ in enumerate(self._freq):
            for j in range(self._time.shape[1]):

                r[j*self._r.shape[1]:(j+1)*self._r.shape[1], i] = self._r[i, :]
        return r

    def compute_r_2D(self):
        r2D = np.zeros((self._freq.shape[0], self._r.shape[1], self._time.shape[1]))
        for f, _ in enumerate(self._freq):
            r2D[f, :, :] = np.repeat(np.transpose(self._r[f, :])[:, np.newaxis], self._time.shape[1], axis=1)
        return r2D

    def reshape_t(self):
        t = np.zeros((self._r.shape[1] * self._time.shape[1], self._freq.shape[0]))
        for i, _ in enumerate(self._freq):
            t[:, i] = np.repeat(self._time[i, :], self._r.shape[1])
        return t