"""Classes for generating causal DAGs and sets of causal DAGs"""

from typing import Tuple, List, Dict, Callable
from scipy.special import comb
import copy
import random
import networkx as nx
from tqdm import tqdm
import pickle
import warnings
from CausalPlayground import StructuralCausalModel


class SCMGenerator:
    """Class for generating SCMs"""

    seed: int
    """Seed for the random number generator."""
    all_functions: Dict[str, Callable]
    """Dictionary of possible functions that can be generated in SCMs generated by this generator. The callable that
    is defined in this dictionary should return the function that is used in the SCM."""
    def __init__(self, all_functions: Dict[str, Callable], seed: int):
        """
        Constructor for the SCM generator.

        :param all_functions: dictionary containing possible functions that can be generated in SCMs generation.
        :param seed: the seed for the random number generator.
        """
        self.seed = seed
        self.all_functions = all_functions

    def create_random(self, possible_functions: List[str], n_endo: int, n_exo: int, exo_distribution: Callable = None,
                      exo_distribution_kwargs: dict = None, allow_exo_confounders: bool = False)\
            -> Tuple[StructuralCausalModel, set]:
        """
        Creates and returns a random StructualCausalModel by first generating a random graph and then randomly choosing
        a function for each edge in the graph based on the possible functions.

        :param possible_functions: list of function that can be used as causal relations. These should correspond to one
        of the strings defined in `self.all_functions`.
        :param n_endo: number of endogenous variables.
        :param n_exo: number of exogenous variables.
        :param exo_distribution: distribution of the exogenous variables. This distribution is applied to all
        exogenous variables.
        :param exo_distribution_kwargs: keyword arguments for the distribution of exogenous variables
        :param allow_exo_distribution: true if exogenous confounders should be generated.

        :return: the random scm and the set of edges that have been removed
        """
        if n_exo > 0:
            assert exo_distribution, 'please provide a distribution for your exogenous variables'
            assert not exo_distribution_kwargs == None, 'please provide kwargs for the distribution of your exogenous variables'
        random.seed(self.seed)
        graph_generator = CausalGraphGenerator(n_endo, n_exo, allow_exo_confounders, seed=self.seed)
        graph, removed_edges = graph_generator.generate_random_graph()
        self.seed += 1
        return self.create_scm_from_graph(graph, possible_functions, exo_distribution, exo_distribution_kwargs),\
               removed_edges

    def create_scm_from_graph(self, graph: nx.DiGraph, possible_functions: List[str], exo_distribution: Callable,
                              exo_distribution_kwargs: dict) -> StructuralCausalModel:
        """
        Generates an SCM following the structure of a given directed graph. A function is randomly chosen from
        `possible_functions` for each edge in the graph.

        :param graph: causal graph
        :param possible_functions: list of function that can be used as causal relations. These should correspond to one
        of the strings defined in 'self.all_functions'
        :param exo_distribution: distribution of the exogenous variables. This distribution is applied to all
        exogenous variables.
        :param exo_distribution_kwargs: keyword arguments for the distribution of exogenous variables
        """
        scm = StructuralCausalModel()
        random.seed(self.seed)

        for n in graph.nodes:
            parents = [p for p in graph.predecessors(n)]
            if graph.nodes[n]['type'] == 'endo':
                # randomly choose one of the possible functions for the current node
                current_function = random.choice(possible_functions)
                scm.add_endogenous_var(n, self.all_functions[current_function](parents), {p: p for p in parents})
            else:
                scm.add_exogenous_var(n, exo_distribution, exo_distribution_kwargs)

        self.seed += 1
        return scm


class CausalGraphGenerator:
    """Class for generating random directed acyclic graphs."""

    allow_exo_confounders: bool
    """flag whether the generated graphs can have exogenous confounders. If False, there will be at most 1 exogenous 
    variable per endogenous variable."""
    seed: int
    """Seed for the random number generator."""
    endo_vars: List[str]
    """List of variable names of the endogenous variables. For the automatic generation in this class, endogenous 
    variables will be named 'Xn' with n ascending from 0."""
    exo_vars: List[str]
    """List of variable names of the exogenous variables. For the automatic generation in this class, exogenous 
    variables will be named 'Un' with n ascending from 0."""
    potential_causes: Dict[str, List[str]]
    """Dictionary of potential causes for each variable in the graph. This helps to manage exogenous variables and their
    confounding effects."""
    fully_connected_graph: nx.DiGraph
    """Fully connected graph of the endogenous variables where every endogenous node is connected to every other 
    endogenous node. The edges from exogenous variables depend on the 'allow_exogenous_confounders' flag."""
    def __init__(self, n_endo: int, n_exo: int = 0, allow_exo_confounders: bool = False, seed: int = 42):
        """
        Constructor defining parameters for the randomly generated graphs.

        :param n_endo: number of endogenous variables.
        :param n_exo: number of exogenous variables.
        :param allow_exo_confounders: flag whether the generated graphs can have exogenous confounders. If False, there
        will be at most 1 exogenous variable per endogenous variable.
        :param seed: seed for the random number generator.
        """
        self.allow_exo_confounders = allow_exo_confounders
        self.seed = seed

        # create var names
        self.endo_vars = ['X' + str(i) for i in range(n_endo)]
        self.exo_vars = ['U' + str(i) for i in range(n_exo)]

        # determine potential causes for each endogenous var
        self.potential_causes = {}
        exo_copy = copy.deepcopy(self.exo_vars)
        for v in self.endo_vars:
            if allow_exo_confounders:  # if confounding is allowed, any node can be the parent of an endogenous var
                self.potential_causes[v] = self.endo_vars + self.exo_vars
            else:  # else, only one exo var can be the cause of an endogenous var
                if not len(exo_copy) == 0:
                    self.potential_causes[v] = self.endo_vars + [exo_copy.pop()]
                else:
                    self.potential_causes[v] = self.endo_vars + []
            # disable self-cycles
            self.potential_causes[v].remove(v)

        del exo_copy
        # create a fully-connected graph based on the confounder settings and store it for later generation
        self.fully_connected_graph = self._make_fully_connected_dag()

    def _make_fully_connected_dag(self) -> nx.DiGraph:
        """
        Helper function that creates and returns a fully connected, possibly cyclic graph. In this graph the exogenous
        variables are roots.

        :return: A fully connected, directed graph.
        """
        graph = nx.DiGraph()
        [graph.add_node(u, type='exo') for u in self.exo_vars]
        [graph.add_node(v, type='endo') for v in self.endo_vars]
        for n, causes in self.potential_causes.items():
            [graph.add_edge(c, n) for c in causes]
        return graph

    def generate_random_graph(self, method: str = None, p: float = 0.3) -> Tuple[nx.DiGraph, set[Tuple[str, str]]]:
        """
        Creates and returns a random nx.DiGraph either by first creating a fully connected graph or creating an Edos
        Reny graph and then randomly deleting one edge after the other until it is acyclic.

        When generating Edos Reny graphs, exogenous variables are not supported for generation. (You can always add
        additional nodes and edges after generation).

        :param method: either 'ER' to generate an Edos Reny graph, or any other string for creating a graph based on
                        random deletion of edges from a fully-connected graph
        :param p: probability of an edge being added to the graph when using the ER method.
        :return: the random nx.DiGraph and the set of edges that have been removed from the fully-connected graph.
        """
        if method == 'ER':  # create erdos reny graph
            assert len(self.exo_vars) == 0, "exogenous variables are not supported when generating ER graphs"
            graph, removed_edges = self._create_er_graph(p)
        else:
            graph, removed_edges = self._create_from_fully_connected()
        self.seed += 1
        return graph, removed_edges

    def _create_er_graph(self, p: float = 0.1) -> (nx.DiGraph, set[Tuple[str, str]]):
        """
        Helper method that creates an Edos Reny graph and then deletes edges until the graph is acyclic.

        :param p: probability of an edge being added to the graph.
        :return: the acyclic nx.DiGraph and the set of edges that have been removed from the fully-connected graph.
        """
        graph = nx.fast_gnp_random_graph(len(self.endo_vars), p=p, seed=self.seed, directed=True)

        # rename the nodes and types according to the fully connected graph
        graph = nx.relabel_nodes(graph, {n: m for (n, m) in zip(graph.nodes, self.fully_connected_graph.nodes)})
        nx.set_node_attributes(graph,
                               {n: {'type': self.fully_connected_graph.nodes()[m]['type']} for (n, m) in
                                zip(graph.nodes, self.fully_connected_graph.nodes)})

        # make acyclic
        while not nx.is_directed_acyclic_graph(graph):
            random_edge = random.choice(list(graph.edges()))
            graph.remove_edge(random_edge[0], random_edge[1])

        # compute which edges are removed
        removed_edges = nx.difference(self.fully_connected_graph, graph).edges

        return graph, set(removed_edges)

    def _create_from_fully_connected(self) -> (nx.DiGraph, set[Tuple[str, str]]):
        """
        Helper method that deletes edges from the fully-connected graph until it is acyclic.

        :return: the acyclic nx.DiGraph and the set of edges that have been removed from the fully-connected graph.
        """
        # generate fully connected graph
        graph = copy.deepcopy(self.fully_connected_graph)
        # delete random edges from the endogenous subgraph until acyclic
        removed_edges = set()
        random.seed(self.seed)
        # determine outgoing edges from exogenous variables
        endogenous_edges = list(graph.edges)
        if not self.allow_exo_confounders:  # only delete the exogenous edges if we allow confounding
            [endogenous_edges.remove(e) for e in graph.edges if e[0] in self.exo_vars]

        # delete edges until the graph is acyclic
        while not nx.is_directed_acyclic_graph(graph):
            random_edge = random.sample(endogenous_edges, 1)[0]
            removed_edges.add(random_edge)
            endogenous_edges.remove(random_edge)
            graph.remove_edge(random_edge[0], random_edge[1])
        return graph, set(removed_edges)


class CausalGraphSetGenerator:
    """
    Represents a collection of unique DAGs and facilitates their bookkeeping.
    """

    generator: CausalGraphGenerator
    """Object generating the causal graphs."""
    graphs: List[nx.DiGraph]
    """List containing all generated graphs."""
    max_endo_dags: int
    """Number of all possible DAGs with 'n_endo' nodes."""
    def __init__(self, n_endo: int, n_exo: int = 0, allow_exo_confounders: bool = False, seed: int = 42):
        """
        Constructor to parametrize the CausalGraphSetGenerator.

        :param n_endo: number of endogenous variables.
        :param n_exo: number of exogenous variables.
        :param allow_exo_confounders: flag whether the generated graphs can have exogenous confounders. If False, there
        will be at most 1 exogenous variable per endogenous variable.
        :param seed: seed for the random number generator.
        """
        self.generator = CausalGraphGenerator(n_endo, n_exo=n_exo,
                                              allow_exo_confounders=allow_exo_confounders,
                                              seed=seed)
        self.graphs = []
        self.max_endo_dags = CausalGraphSetGenerator.max_n_dags(n_endo)

    def generate(self, n: int, method: str = None, p: float = 0.1):
        """
        Generate n distinct causal DAGs. The newly generated graphs are only added to the set of graphs, if the edges
        that have been removed during the generation are different from all other graphs in the set. This ensures the
        uniqueness of the generated, random graphs.

        :param n: how many DAGs to create
        :param method: which method to use to generate the graphs.
        :param p: probability with which an edge is added if ER is used.
        """
        # check whether more DAGs are to be created then combinatorically possible. Only do this if n is not
        # too big because computation takes forever for n > 20 and for such values there exist over 2.3e+72
        # different graphs
        n_exo = len(self.generator.exo_vars)
        if n > self.max_endo_dags * (n_exo + 1):
            n = self.max_endo_dags * (n_exo + 1)
            warnings.warn('Only ', self.max_endo_dags * (n_exo + 1), ' graphs can be created.')

        self.graphs = []
        rem_edges_list = []
        resampled = 0
        print('Creating graphs...')
        pbar = tqdm(total=n - 1)
        while len(self.graphs) < n - 1:
            graph, rem_edges = self.generator.generate_random_graph(method=method, p=p)
            if any([rem_edges == other for other in rem_edges_list]):
                resampled += 1
                continue
            else:
                self.graphs.append(graph)
                rem_edges_list.append(rem_edges)
                pbar.update(1)
        pbar.close()
        print(resampled, 'models resampled')

    def save(self, filepath: str, mode: str = 'wb'):
        """
        Save the generated graphs to a file.

        :param filepath: path to the file.
        :param mode: mode in which to open the file.
        """
        with open(str(filepath), str(mode)) as f:
            pickle.dump(self.graphs, f)

    def load(self, filepath: str):
        """
        Load graphs from files. Warning: does not check whether the loaded graphs are unique.

        :param filepath: path to the file containing the graphs.
        """
        # Warning: doesn't check whether the loaded graphs are actually a set
        with open(str(filepath), 'rb') as f:
            graphs = pickle.load(f)
            if not isinstance(graphs[0], nx.DiGraph):
                raise TypeError('The file does not contain nx.DiGraphs')
            else:
                self.graphs = graphs

    @staticmethod
    def max_n_dags(n_vertices: int) -> int:
        """
        Computes the maximal number of different DAGs over n_vertices nodes. Implemented as in Robinson (1973)

        :param n_vertices: number of nodes in the DAG.
        :return: max number of possible DAGs.
        """
        if n_vertices < 0:
            return 0
        elif n_vertices == 0:
            return 1
        else:
            summ = 0
            for k in range(1, n_vertices + 1):
                summ += (-1) ** (k - 1) * comb(n_vertices, k) * 2 ** (
                        k * (n_vertices - k)) * CausalGraphSetGenerator.max_n_dags(n_vertices - k)
            return int(summ)
