import numpy as np

from f_models import FModel, FBrune, FBoatwright
from SpectralParameters import SpectralParams
from spectral_definitions import damping


class FNormResults:
    def __init__(self, misfit, freq_bins, amplitude_spectrum, amp_theor_spec):
        self.misfit = misfit
        self.freq = freq_bins
        self.amplitude_spectrum = amplitude_spectrum
        self.amplitude_theoretical_spectrum = amp_theor_spec

    def __str__(self):
        return f"Misfit: {self.misfit:.2f}"


class FNorm:
    def __init__(self, norm, spectral_params: SpectralParams, freq_bins, amplitude_spectrum, weights, travel_time,
                 source_model: FModel, logger):
        self.norm = norm
        self.spectral_par = spectral_params
        self.freq = freq_bins
        self.amp_spectrum = amplitude_spectrum
        self.th_amp_sp_damp = None
        self.weights = weights
        self.travel_time = travel_time
        self.f_model = source_model
        self.logger = logger

    def calculate(self):

        self.f_model.sp_par = self.spectral_par
        self.f_model.freq = self.freq

        # theoretical spectrum model (Brune or Boatwright)
        amp_th = self.f_model.fmodel()
        # calculation of quality (damping) values for all frequencies)
        damping_val = damping(q_factor=self.spectral_par.q, frequencies=self.freq, travel_time=self.travel_time)

        if self.norm == "L1":
            misfit = self.f_norm_l1(damping_amp=damping_val, amp_th=amp_th)
        elif self.norm == "L2":
            misfit = self.f_norm_l2(damping_amp=damping_val, amp_th=amp_th)
        else:
            self.logger.error(f"{self.norm} is a wrong norm. It should be L1 or L2")
            return None

        # Correcting theoretical model (Brune or Boatwright) for damping (Q factor)
        self.th_amp_sp_damp = amp_th / damping_val

        return FNormResults(misfit=misfit, freq_bins=self.freq, amplitude_spectrum=self.amp_spectrum,
                            amp_theor_spec=self.th_amp_sp_damp)

    def f_norm_l1(self, damping_amp, amp_th):
        amp_residuals = self.weights * np.abs(np.log10(self.amp_spectrum * damping_amp) - np.log10(amp_th))
        return np.sum(amp_residuals) / np.size(amp_residuals)

    def f_norm_l2(self, damping_amp, amp_th):
        amp_residuals = self.weights * (np.log10(self.amp_spectrum * damping_amp) - np.log10(amp_th)) ** 2
        return np.sqrt(np.sum(amp_residuals) / amp_residuals.size)