import copy
import sys
import time
from typing import Callable, Optional

import libcasm.clusterography as casmclust
import libcasm.configuration as casmconfig
from casm.bset.clexwriter._cpp_str import (
    CppFormatProperties,
)
from casm.bset.cluster_functions import get_occ_site_functions
from casm.bset.polynomial_functions import (
    PolynomialFunction,
)
from libcasm.clexulator import (
    PrimNeighborList,
)

from ._methods import (
    make_neighborhoods,
    make_orbit_bfuncs,
    make_site_bfuncs,
)


class WriterV1Basic:
    """Generates the variables needed for the `v1.basic` Clexulator, the
    standard CASM v1 compatible Clexulator (no automatic differentiation)"""

    @staticmethod
    def template_name():
        """The name of the template file used by this writer

        Returns
        -------
        template_name: str

            The name of the template file used by this writer, as a path relative to
            the casm.bset template directory.
        """
        return "v1.basic/Clexulator.cpp"

    def __init__(
        self,
        prim: casmconfig.Prim,
        project_name: str,
        bset_name: str,
        is_periodic: bool,
        i_clex: Optional[int],
        prim_neighbor_list: PrimNeighborList,
        clusters: list[list[casmclust.Cluster]],
        functions: list[list[list[PolynomialFunction]]],
        occ_site_functions: list[dict],
        occ_site_functions_info: dict,
        linear_function_indices: Optional[set[int]] = None,
        make_variable_name_f: Optional[Callable] = None,
        local_discrete_dof: Optional[dict] = None,
        cpp_fmt: Optional[CppFormatProperties] = None,
        verbose: bool = False,
    ):
        """

        .. rubric:: Constructor

        Parameters
        ----------
        prim: libcasm.configuration.Prim
            The prim.
        project_name: str
            Project name. Used to construct the Clexulator class name. This must
            consist of alphanumeric characters and underscores only. The first
            character may not be a number.
        bset_name: str
            Basis set name. Used to construct the Clexulator class name. This must
            consist of alphanumeric characters and underscores only.
        is_periodic: bool
            Whether the Clexulator is for a periodic basis set or local basis set.
        i_clex: Optional[int]
            If a local basis set, the equivalent clexulator index (start from 0); or
            None to write the prototype.
        prim_neighbor_list: libcasm.clexulator.PrimNeighborList
            The :class:`~libcasm.clexulator.PrimNeighborList` which has been used to
            set the `neighborhood_site_index` for the
            :class:`~casm.bset.polynomial_functions.Variable` used by the `functions`.
        clusters: list[list[libcasm.clusterography.Cluster]]
            The clusters for which cluster functions have been constructed. The cluster
            ``clusters[i_orbit][i_equiv]`` is the `i_equiv`-th symmetrically
            equivalent cluster in the `i_orbit`-th orbit. The order of sites in the
            clusters is not arbitrary, it must be consistent with the
            `cluster_site_index` of the
            :class:`~casm.bset.polynomial_functions.Variable` used in the
            :class:`~casm.bset.polynomial_functions.PolynomialFunction` stored in
            :py:data:`~ClusterFunctionsBuilder.functions`. This can be generated by
            :class:`~casm.bset.cluster_functions.ClusterFunctionsBuilder`.
        functions: list[list[list[casm.bset.polynomial_functions.PolynomialFunction]]]
            List of functions, as generated by
            :class:`~casm.bset.ClusterFunctionsBuilder`, where function
            ``functions[i_orbit][i_cluster][i_func]`` is the `i_func`-th function
            on the `i_cluster`-th cluster in the `i_orbit`-th orbit.
        occ_site_functions: list[dict]
            List of occupation site basis functions. For each sublattice with discrete
            site basis functions, must include:

            - `"sublattice_index"`: int, index of the sublattice
            - `"value"`: list[list[float]], list of the site basis function values,
              as ``value[function_index][occupant_index]``.

        occ_site_functions_info: dict
            Information about occupation site basis functions.

            Occupation site basis functions info, with format:

            - `"max_function_index"`: int, The maximum site function index, across all
              sublattices.
            - `"all_sublattices_have_same_site_functions"`: bool, True if all
              sublattices have same site functions; False otherwise.
            - `"occ_var_name"`: str, A variable name template for the site functions,
              which may be formated using `b` for sublattice index and `m` for site
              function index (i..e ``occ_var_name.format(b=0, m=1)``).
            - `"occ_var_desc"`: str, A description of the occupation
              variable, including a description of the subscript indices.
            - `"occ_var_indices"`: list[list[str, str]], A list of lists, where each
              sublist contains the variable name and description for each subscript
              index.

        linear_function_indices : Optional[set[int]]
            The linear indices of the functions that will be included. If None,
            all functions will be included in the Clexulator. Otherwise,
            only the specified functions will be included in the Clexulator. This
            is an experimental feature.

        make_variable_name_f: Optional[Callable] = None
            Allows specifying a custom class to construct variable names. The default
            class used is :class:`~casm.bset.cluster_functions.MakeVariableName`.
            Custom classes should have the same `__call__` signature as
            :class:`~casm.bset.cluster_functions.MakeVariableName`, and have
            `occ_var_name` and `occ_var_desc` attributes.

        local_discrete_dof: Optional[list[str]] = None
            The types of local discrete degree of freedom (DoF).

        cpp_fmt: Optional[CppFormatProperties] = None,
            C++ string formatting properties. If None, the default values are used.

            .. code-block:: Python

                CppFormatProperties(
                    coeff_fmt_spec=".10f",
                    coeff_atol=1e-10,
                )

        verbose: bool = False
            Print progress statements

        """

        ### defaults ###
        if cpp_fmt is None:
            cpp_fmt = CppFormatProperties()

        ### set variables ###

        # set clexulator_name
        self.clexulator_name = None
        """str: Name of the Clexulator.
        
        This is used for the Clexulator class and must be a valid C++ identifier. It
        is also used for the Clexulator source file name, with extension ".cpp" added.
        The normal convention is that it is:
        
        - For periodic / local prototype: "{project_name}_Clexulator_{bset_name}"
        - For local equivalents: "{project_name}_Clexulator_{bset_name}_{i_clex}"
        
        """

        if is_periodic or i_clex is None:
            self.clexulator_name = f"{project_name}_Clexulator_{bset_name}"
        else:
            self.clexulator_name = f"{project_name}_Clexulator_{bset_name}_{i_clex}"

        self.prim_neighbor_list = prim_neighbor_list
        """libcasm.clexulator.PrimNeighborList: The PrimNeighborList which has been \
        used to set the `neighborhood_site_index` for the 
        :class:`~casm.bset.polynomial_functions.Variable` used by the `functions`."""

        self.occ_site_functions = copy.deepcopy(occ_site_functions)
        """list[dict]: List of occupation site basis functions. 
        
        For each sublattice with discrete site basis functions, includes:

        - `"sublattice_index"`: int, index of the sublattice
        - `"n_occupants"`: int, number of allowed occupants
        - `"value"`: list[list[float]], list of the site basis function values,
          as ``value[function_index][occupant_index]``
        - `"constant_function_index"`: int, index of the constant site basis function
          
        """
        for site_funcs in self.occ_site_functions:
            site_funcs["n_occupants"] = len(site_funcs["value"][0])

            for i in range(len(site_funcs["value"])):
                # note constant occupation functions
                phi = get_occ_site_functions(
                    occ_site_functions=occ_site_functions,
                    sublattice_index=site_funcs["sublattice_index"],
                    site_function_index=i,
                )

                if (phi == 1.0).all():
                    site_funcs["constant_function_index"] = i

            if "constant_function_index" not in site_funcs:
                raise ValueError(
                    f"Occupation site basis functions for sublattice "
                    f"{site_funcs['sublattice_index']} must include "
                    "a constant function."
                )

        self.occ_site_functions_info = occ_site_functions_info
        """dict: Information about occupation site basis functions.
        
        Occupation site basis functions info, with format:
        
        - `"max_function_index"`: int, The maximum site function index, across all
            sublattices.
        - `"all_sublattices_have_same_site_functions"`: bool, True if all _sublattices
            have same site functions; False otherwise.
        - `"occ_var_name"`: str, A variable name template for the site functions,
            which may be formated using `b` for sublattice index and `m` for site 
            function index (i..e ``occ_var_name.format(b=0, m=1)``).
        - `"occ_var_desc"`: str, A description of the occupation
            variable, including a description of the subscript indices.
        - `"occ_var_indices"`: list[list[str, str]], A list of lists, where each sublist
          contains the variable name and description for each subscript index.
        
        """

        ## set linear_function_indices
        self.linear_function_indices = linear_function_indices
        """Optional[set[int]]: The linear indices of the functions that will be 
        included.

        If None, all functions will be included in the Clexulator. Otherwise, only the 
        specified functions will be included in the Clexulator. This
        is an experimental feature."""

        ## set cpp_fmt
        self.cpp_fmt = cpp_fmt
        """libcasm.bset.clexwriter.CppFormatProperties: C++ string formatting \
        properties."""

        ## set n_corr
        self.n_corr = None
        """int: Total number of correlations.

        This is the total number of correlations that could be calculated and does not
        change if `linear_function_indices` is used to prevent evaluation of 
        unnecessary correlations. 
        """

        self.n_corr = 1  # include constant term
        for i_orbit, orbit_functions in enumerate(functions):
            prototype_functions = orbit_functions[0]
            self.n_corr += len(prototype_functions)

        ## set neighborhood info
        if verbose:
            start = time.time()
            print("Generating neighborhoods...")
            sys.stdout.flush()

        _complete_neighborhood, _function_neighborhoods = make_neighborhoods(
            prim=prim,
            is_periodic=is_periodic,
            prim_neighbor_list=prim_neighbor_list,
            clusters=clusters,
            functions=functions,
            linear_function_indices=linear_function_indices,
            make_variable_name_f=make_variable_name_f,
            local_discrete_dof=local_discrete_dof,
        )

        if verbose:
            print("Generating neighborhoods DONE")
            elapsed_time = time.time() - start
            print(f"time: {elapsed_time:0.4f} (s)")
            print()
            sys.stdout.flush()

        self.complete_neighborhood = _complete_neighborhood
        """dict: The neighborhood needed to evaluate basis functions 
        involving sites in the origin unit cell. 
        
        Contains:
        
        - ``"unitcells"``: list[list[int]], List of ``[i,j,k]`` unit
          cell indices for the neighborhood.
        - ``"sites"``: list[list[int]], List of ``[b,i,j,k]`` site
          indices for the neighborhood.
        
        """

        self.function_neighborhoods = _function_neighborhoods
        """list[dict]: The neighborhood data for each orbit of basis functions. 
        
        Each dictionary contains:

        - ``"linear_function_index"``: int, Linear function index
        - ``"same_as"``: Optional[int], If not None, the linear function index of the
          first orbit of functions in the same cluster orbit as this function, which
          has the same neighborhood.
        - ``"unitcells"``: list[list[int]], List of ``[i,j,k]`` unit
          cell indices for the neighborhood of the orbit of basis function. This is
          only populated for the first function in each cluster orbit.
        - ``"sites"``: list[list[int]], List of ``[b,i,j,k]`` site
          indices for the neighborhood of the orbit of basis function. This is
          only populated for the first function in each cluster orbit.
        
        """

        ## set nlist_size
        import libcasm.xtal as xtal

        nlist_size = 0
        for bijk in self.complete_neighborhood["sites"]:
            site = xtal.IntegralSiteCoordinate.from_list(bijk)
            nlist_size = max(nlist_size, prim_neighbor_list.neighbor_index(site) + 1)
        self.nlist_size = nlist_size
        """int: Maximum neighbor list index of sites involved in evaluating basis 
        functions."""

        ## set orbit_bfuncs and orbit_bfuncs_variables_needed
        self.orbit_bfuncs = None
        """list[dict]: Data for each orbit basis function (i.e. if periodic, the 
        contribution to the global correlations from one unit cell for periodic; if
        local, the sum of all symmetrically equivalent local cluster functions). 
        
        Includes:

        - ``"linear_function_index"``: int, Linear function index
        - ``"linear_orbit_index"``: int, Linear cluster orbit index
        - ``"cpp"``: str, C++ expression for evaluating the global correlation
          contribution
        - ``"latex_prototype"``: str, Latex formula for the prototype cluster
        - ``"latex_orbit"``: str, Latex formula for the orbit contribution
        
        """

        self.orbit_bfuncs_variables_needed = None
        """dict[str,list[list[int]]]: The variables needed to evaluate the orbit basis 
        functions. 
        
        The dictionary keys are DoF type key, and values are lists of
        [`component_index`, `neighbor_list_index`, `sublattice_index`] for variables
        of that type. The `neighbor_list_index` and `sublattice_index` for global
        continuous DoF are None.
        """

        if verbose:
            start = time.time()
            print("Generating orbit function formulas...")
            sys.stdout.flush()

        self.orbit_bfuncs, self.orbit_bfuncs_variables_needed = make_orbit_bfuncs(
            prim_neighbor_list=self.prim_neighbor_list,
            clusters=clusters,
            functions=functions,
            occ_site_functions=self.occ_site_functions,
            cpp_fmt=self.cpp_fmt,
            linear_function_indices=self.linear_function_indices,
        )

        if verbose:
            print("Generating orbit function formulas DONE")
            elapsed_time = time.time() - start
            print(f"time: {elapsed_time:0.4f} (s)")
            print()
            sys.stdout.flush()

        ## set site_bfuncs and site_bfuncs_variables_needed_at

        self.site_bfuncs = None
        """list[dict]: List of dictionaries, one per cluster expansion basis function, 
        with information needed for printing the Clexulator point correlation 
        evaluation methods.
        
        The format is:

        - "linear_function_index": int - The cluster expansion basis function index.
        - "linear_orbit_index": int - The cluster orbit index
        - "at": list[dict] - Data for functions at specific points

          - "neighbor_list_index": int - Index into the neighbor list of the point
          - "cpp": str - The C++ code to evaluate the point functions
          - "occ_delta_cpp": str - The C++ code to evaluate the change in point
            functions due to the change of occupation index on the site from
            ``"occ_i"`` to ``"occ_f"``.
          - "latex_prototype": str - The latex expression for the point functions,
            including only functions from the prototype cluster
          - "latex_orbit": str - The latex expression for the point functions,
            including functions from the entire cluster orbit
        
        """

        self.site_bfuncs_variables_needed_at = None
        """list[dict[str,list[list[int]]]]: For each point, the variables needed to 
        evaluate the point functions. 
        
        The value of `site_bfuncs_variables_needed_at[neighbor_list_index]` is a 
        dictionary, where the keys are DoF type key, and values are lists of
        [`component_index`, `neighbor_list_index`, `sublattice_index`] for variables
        of that type. The `neighbor_list_index` and `sublattice_index` for global
        continuous DoF are None.
        """

        if verbose:
            start = time.time()
            print("Generating site function formulas...")
            sys.stdout.flush()

        self.site_bfuncs, self.site_bfuncs_variables_needed_at = make_site_bfuncs(
            prim=prim,
            is_periodic=is_periodic,
            prim_neighbor_list=self.prim_neighbor_list,
            clusters=clusters,
            functions=functions,
            occ_site_functions=self.occ_site_functions,
            cpp_fmt=self.cpp_fmt,
            linear_function_indices=self.linear_function_indices,
            make_variable_name_f=make_variable_name_f,
            local_discrete_dof=local_discrete_dof,
        )

        if verbose:
            print("Generating site function formulas DONE")
            elapsed_time = time.time() - start
            print(f"time: {elapsed_time:0.4f} (s)")
            print()
            sys.stdout.flush()

        ## set n_point_corr_sites
        n_point_corr_sites = 0
        for x in self.site_bfuncs:
            for y in x["at"]:
                if y["n_point_functions"] > 0:
                    n_point_corr_sites = max(
                        n_point_corr_sites, y["neighbor_list_index"] + 1
                    )
        self.n_point_corr_sites = n_point_corr_sites
        """int: Number of sites where point correlations can be evaluated.

        This is the maximum neighbor list index for sites that have site functions.
        For periodic cluster expansions, the maximum value this can take is the number 
        of sublattices included in the `prim_neighbor_list`. For local cluster 
        expansions, the maximum value this can take is the maximum neighbor list index
        of the complete site neighborhood. It is possible that some of the sites do not 
        actually have any site functions associated with them, in which case the point 
        correlations all evaluate as zero.
        """

        # if is_periodic:
        #     self.n_point_corr_sites = len(prim_neighbor_list.sublattice_indices())
        # else:
        #     self.n_point_corr_sites = prim_neighbor_list.n_neighborhood_sites()

        ## set continuous DoF and params (parameter pack implementation details)

        self.continuous_dof = []
        """list[dict]: List of continuous DoF to register.

        For each, must include:

        - `"key"`: str, DoF name (i.e. ``"Hstrain``)
        - `"is_global"`: bool, DoF type
        - `"max_n_components"`: int, Maximum number of components (for local DoF)
        - `"sites"`: list[dict], Describes local DoF, includes:

          - `"sublattice_index"`: int, Sublattice index containing local DoF
          - `"n_components"`: int, Number of components on this sublattice

        """

        self.params = []
        """list[dict]: List of parameter packs to include.

        Parameter packs are a data structure used internally by the Clexulator that
        store the evaluated variables used by the cluster functions. The data structure
        type can be written in a generic fashion to allow automatic differentiation
        when it is enabled. There should be one parameter pack for each DoF included 
        in the basis functions and one for correlations. For each parameter pack, must 
        include:

        - `"name"`: str, Parameter name (i.e. ``"occ_site_func"``,
          ``"Hstrain_var"``, ``"disp_var"``, or ``"corr"``.)
        - `"rows"`: int, Number of rows needed by the parameter pack. For DoF,
          this is the maximum number of components. For "corr", this is the number
          of functions.
        - `"cols"`: int, Number of columns needed by the parameter pack. For site
          DoF this is the number of neighbors, for global DoF this is 1,
          and for "corr" this is 1.
        - `"is_independent"`: str, String representation of C++ bool. For DoF
          this is ``"true"``, and for "corr" this is ``"false"``.

        """

        # add correlations parameter pack
        corr_param = {
            "name": "corr",
            "rows": self.n_corr,
            "cols": 1,
            "is_independent": "false",
        }
        self.params.append(corr_param)

        # add dof parameter packs and continuous dof info
        for key, variables_needed in self.orbit_bfuncs_variables_needed.items():
            if len(variables_needed) == 0:
                continue

            # check variables of type==key
            is_global = True
            max_component_index = 0
            max_component_index_by_site = {}
            for var in variables_needed:
                component_index = var[0]
                sublattice_index = var[2]
                if sublattice_index is not None:
                    is_global = False
                if component_index > max_component_index:
                    max_component_index = component_index
                if sublattice_index is not None:
                    if sublattice_index not in max_component_index_by_site:
                        max_component_index_by_site[sublattice_index] = 0
                    if component_index > max_component_index_by_site[sublattice_index]:
                        max_component_index_by_site[sublattice_index] = component_index
            rows = max_component_index + 1

            if is_global:
                # global Dof only need one column in parameter pack
                cols = 1
                sites_info = None
            else:
                # local Dof need one column per neighborhood site in parameter pack
                cols = self.prim_neighbor_list.n_neighborhood_sites()

                # get max number of components by sublattice
                sites_info = []
                for b, c_max in max_component_index_by_site.items():
                    sites_info.append(
                        {
                            "sublattice_index": b,
                            "n_components": c_max + 1,
                        }
                    )

            if key == "occ":
                is_discrete = True
                name = "occ_site_func"
            else:
                is_discrete = False
                name = f"{key}_var"

            dof_param = {
                "name": name,
                "rows": rows,
                "cols": cols,
                "is_independent": "true",
            }
            self.params.append(dof_param)

            if not is_discrete:
                dof_info = {
                    "key": key,
                    "is_global": is_global,
                }
                if not is_global:
                    dof_info["max_n_components"] = rows
                    dof_info["sites"] = sites_info
                self.continuous_dof.append(dof_info)

        ## set cpp_fmt
        self.cpp_fmt = cpp_fmt
        """casm.bset.clexwriter.CppFormatProperties: C++ string formatting properties.
        """

    def variables(self) -> dict:
        """Returns a dictionary with variables to be rendered using a Jinja2 template.

        Returns
        -------
        variables: dict
            Variables used by the v1.basic/Clexulator.cpp Jinja2 Clexulator template

        """
        return {
            "clexulator_name": self.clexulator_name,
            "nlist_size": self.nlist_size,
            "n_corr": self.n_corr,
            "n_point_corr_sites": self.n_point_corr_sites,
            "orbit_bfuncs": self.orbit_bfuncs,
            "orbit_bfuncs_variables_needed": self.orbit_bfuncs_variables_needed,
            "site_bfuncs": self.site_bfuncs,
            "site_bfuncs_variables_needed_at": self.site_bfuncs_variables_needed_at,
            "occ_site_functions": self.occ_site_functions,
            "occ_site_functions_info": self.occ_site_functions_info,
            "continuous_dof": self.continuous_dof,
            "params": self.params,
            "nlist_weight_matrix": self.prim_neighbor_list.weight_matrix().tolist(),
            "nlist_sublattice_indices": self.prim_neighbor_list.sublattice_indices(),
            "nlist_total_n_sublattice": self.prim_neighbor_list.total_n_sublattice(),
            "complete_neighborhood": self.complete_neighborhood,
            "function_neighborhoods": self.function_neighborhoods,
        }
