"""Module concerned with probability calculation
"""
from __future__ import annotations
import numpy as np
from rich.console import Console


CONSOLE = Console()

class Probability():
    """Class that stores the conditional probability according to frequency intervals
    """
    def __init__(self, freq_min, freq_max, p_periodic = 0, p_freq = 0, p_freq_given_periodic = 0, p_periodic_given_freq=1):
        """init function

        Args:
            freq_min (float): minimum frequencies in group
            freq_max (float): max frequencies in group
            p_periodic (float, optional): probability of being periodic. Defaults to 0.
            p_freq (float, optional): Probability that the signal has the frequency in the range [freq_min, freq_max]. Defaults to 0.
            p_freq_given_periodic (float, optional): conditional probability that expresses how probable the frequency is in the interval [freq_min, freq_max] ,
            given that the signal is periodic. Defaults to 0.
            p_periodic_given_freq (float, optional): Probability that the signal is periodic given that it has a frequency A. Defaults to 1.
        """
        self.freq_min = freq_min
        self.freq_max = freq_max
        self.p_periodic = p_periodic
        self.p_freq =  p_freq
        self.p_freq_given_periodic = p_freq_given_periodic
        self.p_periodic_given_freq = p_periodic_given_freq

    def set(self, p_periodic = np.nan, p_freq = np.nan, p_freq_given_periodic = np.nan, p_periodic_given_freq=np.nan):
        self.p_periodic = p_periodic if not np.isnan(p_periodic) else self.p_periodic
        self.p_freq =  p_freq if not np.isnan(p_freq) else self.p_freq
        self.p_freq_given_periodic = p_freq_given_periodic if not np.isnan(p_freq_given_periodic) else self.p_freq_given_periodic
        self.p_periodic_given_freq = p_periodic_given_freq if not np.isnan(p_periodic_given_freq) else self.p_periodic_given_freq

    def display(self):
        CONSOLE.print(
                    f"[purple][PREDICTOR]:[/] P([{self.freq_min:.3f},{self.freq_max:.3f}] Hz) = {self.p_periodic*100:.2f}%\n"
                    f"[purple][PREDICTOR]:[/] |-> [{self.freq_min:.3f},{self.freq_max:.3f}] Hz = [{1/self.freq_max if self.freq_max != 0 else np.NaN:.3f},{1/self.freq_min if self.freq_min != 0 else np.NaN:.3f}] sec\n"
                    f"[purple][PREDICTOR]:[/] '-> P([{self.freq_min:.3f},{self.freq_max:.3f}] Hz | periodic) = {self.p_freq_given_periodic*100:.2f}%\n"
                )

    def get_freq_prob(self, freq):
        if freq >= self.freq_min and freq <= self.freq_max:
            return True
        else: 
            return False


