from __future__ import annotations

from typing import Optional, Union

import networkx as nx
import numpy as np
from numpy.typing import NDArray

from ..network_generator import NetworkGenerator


class CNVMParameters:
    """
    Container for the parameters of the Continuous-time Noisy Voter Model (CNVM).

    A node i transitions from its current opinion m to a different opinion n at rate
    r[m,n] * d(i,n) / (d(i)^alpha) + r_tilde[m,n],
    where d(i,n) is the count of opinion n in the neighborhood of agent i, and d(i) the degree of node i.

    Either a network has to specified, or a NetworkGenerator,
    or num_agents, in which case a complete network is used.
    If multiple are given, NetworkGenerator overrules network, and network overrules num_agents.

    The rate parameters r and r_tilde can be given as numpy arrays of shape (num_opinions, num_opinions),
    or as floats, in which case all rates are set to this value.

    (Internally, the CNVM uses an equivalent set of rate parameters: r_imit, r_noise, prob_imit and prob_noise.
    A node i transitions from its current opinion m to a different opinion n at rate
    r_imit * d(i,n) / (d(i)^alpha) * prob_imit[m, n] + r_noise * (1/num_opinions) * prob_noise[m, n].
    These parameters can be provided instead of the usual r and r_tilde.)
    """

    def __init__(
        self,
        num_opinions: int,
        num_agents: Optional[int] = None,
        network: Optional[nx.Graph] = None,
        network_generator: Optional[NetworkGenerator] = None,
        alpha: float = 1.0,
        # rate parameters in style 1
        r: Union[float, NDArray, None] = None,
        r_tilde: Union[float, NDArray, None] = None,
        # rate parameters in style 2
        r_imit: Optional[float] = None,
        r_noise: Optional[float] = None,
        prob_imit: Union[float, NDArray] = 1,
        prob_noise: Union[float, NDArray] = 1,
    ):
        self.num_opinions = num_opinions
        self.alpha = alpha

        # network
        self.num_agents, self.network, self.network_generator = _sanitize_network_input(
            num_agents, network, network_generator
        )

        (
            self.r,
            self.r_tilde,
            self.r_imit,
            self.r_noise,
            self.prob_imit,
            self.prob_noise,
        ) = _sanitize_rates_input(
            num_opinions, r, r_tilde, r_imit, r_noise, prob_imit, prob_noise
        )

    def get_network(self) -> nx.Graph:
        """
        If self.network exists, returns the network.
        Else, if a NetworkGenerator was given, returns a network generated by it.
        Else, returns the complete graph.

        Returns
        -------
        nx.Graph
        """
        if self.network is not None:
            return self.network
        elif self.network_generator is not None:
            return self.network_generator()
        else:
            return nx.complete_graph(self.num_agents)

    def set_network(self, network: nx.Graph) -> None:
        """
        Set new network.

        Only works if no NetworkGenerator was given.

        Parameters
        ----------
        network : nx.Graph
        """
        if self.network_generator is not None:
            raise ValueError("Cannot set network when there is a NetworkGenerator.")
        self.network = network
        self.num_agents = len(network.nodes)

    def update_network_by_generator(self) -> None:
        """
        Update self.network via the NetworkGenerator.

        Only works if a NetworkGenerator was given.
        """
        if self.network_generator is None:
            raise ValueError("No NetworkGenerator was given.")
        self.network = self.network_generator()

    def change_rates(
        self,
        r: Union[float, NDArray, None] = None,
        r_tilde: Union[float, NDArray, None] = None,
    ):
        """
        Change one or both rate parameters.

        If only one argument is given, the other rate parameter stays the same.

        Parameters
        ----------
        r : float | np.ndarray, optional
        r_tilde : float | np.ndarray
        """
        r = r if r is not None else self.r
        r_tilde = r_tilde if r_tilde is not None else self.r_tilde
        (
            self.r,
            self.r_tilde,
            self.r_imit,
            self.r_noise,
            self.prob_imit,
            self.prob_noise,
        ) = _sanitize_rates_input(self.num_opinions, r, r_tilde, None, None, 1, 1)


def _sanitize_network_input(
    num_agents: Optional[int],
    network: Optional[nx.Graph],
    network_generator: Optional[NetworkGenerator],
) -> tuple[int, Optional[nx.Graph], Optional[NetworkGenerator]]:
    if network_generator is not None:
        return (network_generator.num_agents, None, network_generator)
    if network is not None:
        return (network.number_of_nodes(), network, None)
    if num_agents is not None:
        return (num_agents, None, None)
    raise ValueError(
        "Either a network or a NetworkGenerator or num_agents has to be specified."
    )


def _sanitize_rates_input(
    num_opinions: int,
    r: Union[float, NDArray, None],
    r_tilde: Union[float, NDArray, None],
    r_imit: Optional[float],
    r_noise: Optional[float],
    prob_imit: Union[float, NDArray],
    prob_noise: Union[float, NDArray],
) -> tuple[NDArray, NDArray, float, float, NDArray, NDArray]:
    if r is not None and r_tilde is not None:  # style 1
        r = _to_matrix(num_opinions, r)
        r_tilde = _to_matrix(num_opinions, r_tilde)
        if np.min(r) < 0 or np.min(r_tilde) < 0:
            raise ValueError("Rates have to be non-negative.")

        (
            r_imit,
            r_noise,
            prob_imit,
            prob_noise,
        ) = convert_rate_to_cnvm(r, r_tilde)

        return (r, r_tilde, r_imit, r_noise, prob_imit, prob_noise)

    if r_imit is not None and r_noise is not None:  # style 2
        prob_imit = _to_matrix(num_opinions, prob_imit)
        prob_noise = _to_matrix(num_opinions, prob_noise)
        if r_imit < 0 or r_noise < 0:
            raise ValueError("Rates have to be non-negative.")
        if np.min(prob_imit) < 0 or np.max(prob_imit) > 1:
            raise ValueError("Probabilities have to be between 0 and 1.")
        if np.min(prob_noise) < 0 or np.max(prob_noise) > 1:
            raise ValueError("Probabilities have to be between 0 and 1.")

        r = r_imit * prob_imit
        r_tilde = r_noise * prob_noise / num_opinions

        return (r, r_tilde, r_imit, r_noise, prob_imit, prob_noise)

    raise ValueError("Rate parameters have to be provided.")


def _to_matrix(num_opinions: int, r: Union[float, NDArray]) -> NDArray:
    """
    If `r` is a already a NDArray, return `r` but with 0 put on the diagonal.
    If `r` is a float, return a `num_opinions` x `num_opinions` matrix with 0 on the diagonal and `r` everywhere else.
    """
    if isinstance(r, (int, float)):
        r = r * np.ones((num_opinions, num_opinions))
    np.fill_diagonal(r, 0)
    return r


def save_params_as_txt_file(filename: str, params: CNVMParameters):
    """
    Save parameters as a readable .txt file.

    Parameters
    ----------
    filename : str
    params : CNVMParameters
    """
    if filename[-4:] != ".txt":
        this_filename = filename + ".txt"
    else:
        this_filename = filename

    with open(this_filename, "w") as f:
        f.write(f"num_agents = {params.num_agents}\n")
        f.write(f"num_opinions = {params.num_opinions}\n")
        f.write(f"alpha = {params.alpha}\n")
        f.write(f"r_imit = {params.r_imit}\n")
        f.write(f"r_noise = {params.r_noise}\n\n")
        f.write(f"prob_imit =\n {params.prob_imit}\n\n")
        f.write(f"prob_noise =\n {params.prob_noise}\n\n")

        if params.network_generator is not None:
            f.write(f"network_generator = {params.network_generator}\n")
        elif params.network is not None:
            f.write(f"network = {params.network}\n")
        else:
            f.write(f"network = fully connected\n")


def convert_rate_to_cnvm(
    r: NDArray, r_tilde: NDArray
) -> tuple[float, float, NDArray, NDArray]:
    """
    Convert the rates r and r_tilde to the parameters used in the CNVM, i.e., r_imit, r_noise, prob_imit, prob_noise.

    The rates r and r_tilde are defined as:
    An agent i transitions from his current opinion m to a different opinion n at rate
    r[m, n] * d(i,n) / (d(i)^alpha) + r_tilde[m, n]
    where d(i,n) is the count of opinion n in the neighborhood of agent i, and d(i) the degree of node i.

    Parameters
    ----------
    r : np.ndarray
        shape = (num_opinions, num_opinions)
    r_tilde : np.ndarray
        shape = (num_opinions, num_opinions)

    Returns
    -------
    tuple[float, float, np.ndarray, np.ndarray]
        r_imit, r_noise, prob_imit, prob_noise
    """
    num_opinions = r.shape[0]

    # r[m,n] = r_imit * prob_imit[m,n]
    this_r = np.copy(r)
    np.fill_diagonal(this_r, 0)
    r_imit = np.max(this_r)
    if r_imit > 0:
        prob_imit = this_r / r_imit
    else:
        prob_imit = np.zeros((num_opinions, num_opinions))

    # r_tilde[m, n] = r_noise / num_opinions * prob_noise[m, n]
    this_r_tilde = np.copy(r_tilde)
    np.fill_diagonal(this_r_tilde, 0)
    r_noise = np.max(this_r_tilde) * num_opinions
    if r_noise > 0:
        prob_noise = this_r_tilde * num_opinions / r_noise
    else:
        prob_noise = np.zeros((num_opinions, num_opinions))

    return r_imit, r_noise, prob_imit, prob_noise


def convert_rate_from_cnvm(params: CNVMParameters) -> tuple[NDArray, NDArray]:
    """
    Convert the rates used in the CNVM to r and r_tilde.

    The rates r and r_tilde are defined as:
    An agent i transitions from his current opinion m to a different opinion n at rate
    r[m, n] * d(i,n) / (d(i)^alpha) + r_tilde[m, n]
    where d(i,n) is the count of opinion n in the neighborhood of agent i, and d(i) the degree of node i.

    Parameters
    ----------
    params : CNVMParameters

    Returns
    -------
    tuple[np.ndarray, np.ndarray]
        r, r_tilde
    """
    r = params.r_imit * params.prob_imit
    r_tilde = params.r_noise * params.prob_noise / params.num_opinions
    return r, r_tilde
