import pathlib
import re
import sys
import time
from typing import Callable, Optional

import jinja2
from sortedcontainers import SortedList

import libcasm.clusterography as casmclust
import libcasm.configuration as casmconfig
import libcasm.xtal as xtal
from casm.bset.clexwriter import (
    CppFormatProperties,
    occ_delta_site_bfunc_cpp_str,
    orbit_bfunc_cpp_str,
    site_bfunc_cpp_str,
)
from casm.bset.cluster_functions import (
    ClexBasisSpecs,
    ClusterFunctionsBuilder,
    get_occ_site_functions,
    make_local_point_functions,
    make_point_functions,
)
from casm.bset.json_io import dump
from casm.bset.polynomial_functions import (
    PolynomialFunction,
)
from libcasm.clexulator import (
    PrimNeighborList,
)


def make_neighborhoods(
    prim: casmconfig.Prim,
    is_periodic: bool,
    prim_neighbor_list: PrimNeighborList,
    clusters: list[list[casmclust.Cluster]],
    functions: list[list[list[PolynomialFunction]]],
    linear_function_indices: Optional[set[int]] = None,
    make_variable_name_f: Optional[Callable] = None,
    local_discrete_dof: Optional[dict] = None,
):
    """Make the neighborhood lists for the functions

    Parameters
    ----------
    prim: libcasm.configuration.Prim
        The prim.

    is_periodic: bool
        Whether the Clexulator is for a periodic basis set or local basis set.

    prim_neighbor_list: libcasm.clexulator.PrimNeighborList
        :class:`~libcasm.clexulator.PrimNeighborList`, containing all the neighbors
        needed for evaluation, as generated by
        :class:`~casm.bset.cluster_functions.ClusterFunctionBuilder`.

    clusters: list[list[libcasm.clusterography.Cluster]]
        List of clusters, as generated by
        :class:`~casm.bset.cluster_functions.ClusterFunctionBuilder`. Site order
        must be consistent with `functions`.

    functions: list[list[list[casm.bset.polynomial_functions.PolynomialFunction]]]
        List of functions, as generated by
        :class:`~casm.bset.cluster_functions.ClusterFunctionBuilder`, where
        ``functions[i_orbit][i_equiv][i_func]``,
        is the `i_func`-th function on the cluster given by
        `clusters[i_orbit][i_equiv]`.

    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

    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).

    Returns
    -------
    complete_neighborhood: dict
        The complete neighborhood of the basis functions. Contains:

        - ``"unitcells"``: list[list[int]], List of ``[i,j,k]`` unit
          cell indices for the neighborhood of the orbit of basis function.
        - ``"sites"``: list[list[int]], List of ``[b,i,j,k]`` site
          indices for the neighborhood of the orbit of basis function.

    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.
    """

    # get total number of basis functions
    n_corr = 1  # include constant term
    for i_orbit, orbit_functions in enumerate(functions):
        prototype_functions = orbit_functions[0]
        n_corr += len(prototype_functions)

    # make neighborhood index -> site lookup
    neighborhood_index_to_site = []
    for unitcell in prim_neighbor_list:
        for sublattice_index in prim_neighbor_list.sublattice_indices():
            integral_site_coordinate = xtal.IntegralSiteCoordinate(
                sublattice=sublattice_index,
                unitcell=unitcell,
            )
            neighborhood_index_to_site.append(integral_site_coordinate)

    function_neighborhoods = []
    for i_func in range(n_corr):
        function_neighborhoods.append(
            {
                "linear_function_index": i_func,
                "same_as": None,
                "unitcells": SortedList(),
                "sites": SortedList(),
            }
        )

    # the constant function is not included in `functions`
    linear_function_index_sum = 1
    for i_orbit, orbit_functions in enumerate(functions):
        # This is orbit_functions[i_equiv][i_func]
        # We want: orbit_equiv_functions[i_func][j], where i_func indicates
        # symmetrically equivalent functions and j is an index over all
        # symmetrically equivalent functions associated with the origin unit cell

        # For periodic:
        # - Get all functions that include sites in the origin unit cell (i.e. add
        #   translational equivalents if they involve sites in the origin unit cell)
        # - Then organize by i_func on prototype -> all equivalent functions
        # For local:
        # - Just re-organize by i_func on prototype -> all equivalent functions

        if is_periodic:
            # point_functions: list[list[list[PolynomialFunction]]]
            #     All point functions for this cluster orbit, for each point corr site
            #     point_functions[i_func][nlist_index][i_point_function]
            point_functions = make_point_functions(
                prim=prim,
                prim_neighbor_list=prim_neighbor_list,
                orbit=clusters[i_orbit],
                orbit_functions=orbit_functions,
                make_variable_name_f=make_variable_name_f,
                local_discrete_dof=local_discrete_dof,
            )

            # Note: point_functions contains duplicates, but that's OK for generating
            # the neighborhood

            # orbit_equiv_functions: list[list[PolynomialFunction]]
            #     All functions for this cluster orbit,
            #     orbit_equiv_functions[i_func][i_equiv_function]
            orbit_equiv_functions = []
            for i_func, functions_by_nbor in enumerate(point_functions):
                x = []
                for funcs in functions_by_nbor:
                    x += funcs
                orbit_equiv_functions.append(x)

        else:
            # orbit_equiv_functions: list[list[PolynomialFunction]]
            #     All point functions for this cluster orbit,
            #     orbit_equiv_functions[i_func][i_equiv_function]
            prototype_basis_set = orbit_functions[0]
            orbit_equiv_functions = [[]] * len(prototype_basis_set)

            for i_equiv, equiv_functions in enumerate(orbit_functions):
                for i_func, f in enumerate(equiv_functions):
                    orbit_equiv_functions[i_func].append(f)

        for i_func, equiv_functions in enumerate(orbit_equiv_functions):
            linear_function_index = linear_function_index_sum + i_func

            if (
                linear_function_indices is not None
                and linear_function_index not in linear_function_indices
            ):
                continue

            if i_func != 0:
                function_neighborhoods[linear_function_index][
                    "same_as"
                ] = linear_function_index_sum
                continue

            for func in equiv_functions:
                for var in func.variables:
                    # skip global variables:
                    if var.neighborhood_site_index is None:
                        continue

                    site = neighborhood_index_to_site[var.neighborhood_site_index]

                    site_as_list = site.to_list()
                    unitcell_as_list = site.unitcell().tolist()

                    f_neighborhood = function_neighborhoods[linear_function_index]
                    f_unitcell_nbors = f_neighborhood["unitcells"]
                    f_site_nbors = f_neighborhood["sites"]

                    if unitcell_as_list not in f_unitcell_nbors:
                        f_unitcell_nbors.add(unitcell_as_list)

                    if site_as_list not in f_site_nbors:
                        f_site_nbors.add(site_as_list)

        linear_function_index_sum += len(orbit_functions[0])

    complete_neighborhood = {
        "unitcells": SortedList(),
        "sites": SortedList(),
    }
    for d in function_neighborhoods:
        d["unitcells"] = list(d["unitcells"])
        d["sites"] = list(d["sites"])

        for unitcell in d["unitcells"]:
            if unitcell not in complete_neighborhood["unitcells"]:
                complete_neighborhood["unitcells"].add(unitcell)

        for site in d["sites"]:
            if site not in complete_neighborhood["sites"]:
                complete_neighborhood["sites"].add(site)
    complete_neighborhood["unitcells"] = list(complete_neighborhood["unitcells"])
    complete_neighborhood["sites"] = list(complete_neighborhood["sites"])

    return (complete_neighborhood, function_neighborhoods)


def _add_variables_needed(
    data: dict[str, list],
    functions: list[PolynomialFunction],
    prim_neighbor_list: PrimNeighborList,
    occ_site_functions: list[list[float]],
):
    nlist_sublat_indices = prim_neighbor_list.sublattice_indices()
    for f in functions:
        for var in f.variables:
            if var.key not in data:
                data[var.key] = []
            # continuous DoF
            n = var.neighborhood_site_index
            if n is None:
                b = None
            else:
                i = var.neighborhood_site_index % len(nlist_sublat_indices)
                b = nlist_sublat_indices[i]
            if var.key == "occ":
                # skip constant occupation functions
                phi = get_occ_site_functions(
                    occ_site_functions=occ_site_functions,
                    sublattice_index=b,
                    site_function_index=var.site_basis_function_index,
                )

                if (phi == 1.0).all():
                    continue

                # occ DoF / site basis function
                value = [var.site_basis_function_index, n, b]
            else:
                value = [var.component_index, n, b]
            if value not in data[var.key]:
                data[var.key].append(value)


def make_orbit_bfuncs(
    prim_neighbor_list: PrimNeighborList,
    clusters: list[list[casmclust.Cluster]],
    functions: list[list[list[PolynomialFunction]]],
    occ_site_functions: list[dict],
    cpp_fmt: CppFormatProperties,
    linear_function_indices: Optional[set[int]] = None,
) -> tuple[list[dict], dict[str, list[list[int]]]]:
    """Convert clusters and functions to data used by Jinja templates to write the \
    orbit basis function evaluating methods (i.e. the contributions to the global \
    correlations from one unit cell).

    Parameters
    ----------
    prim_neighbor_list: libcasm.clexulator.PrimNeighborList
        :class:`~libcasm.clexulator.PrimNeighborList`, containing all the neighbors
        needed for evaluation, as generated by
        :class:`~casm.bset.cluster_functions.ClusterFunctionBuilder`.

    clusters: list[list[libcasm.clusterography.Cluster]]
        List of clusters, as generated by
        :class:`~casm.bset.cluster_functions.ClusterFunctionBuilder`. Site order
        must be consistent with `functions`.

    functions: list[list[list[casm.bset.polynomial_functions.PolynomialFunction]]]
        List of functions, as generated by
        :class:`~casm.bset.cluster_functions.ClusterFunctionBuilder`.

    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 = functions[function_index][occupant_index]``.

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

    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


    Returns
    -------
    orbit_bfuncs: list[dict]
        Data for each orbit basis function (i.e. the contribution to the global
        correlations from one unit cell). 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
        - ``""latex_prototype_with_neighbor_indices": str, Latex formula for the
          prototype using neighbor site indices instead of cluster site indices
        - ``latex_orbit"``: str, Latex formula for the orbit contribution

    variables_needed: 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.

    """
    linear_function_index = 0
    linear_orbit_index = 0
    orbit_bfuncs = []
    variables_needed = {}

    # For constant term, linear_function_index==0:
    if (
        linear_function_indices is None
        or linear_function_index in linear_function_indices
    ):
        orbit_bfuncs.append(
            {
                "linear_function_index": linear_function_index,
                "linear_orbit_index": linear_orbit_index,
                "cpp": "1",
                "latex_prototype": "1",
                "latex_prototype_with_neighbor_indices": "1",
                "latex_orbit": "1",
            }
        )
    linear_function_index += 1

    ### For all cluster orbits (includes null cluster):
    # i_orbit: cluster orbit index
    # orbit: list[Cluster]
    for i_orbit, orbit in enumerate(clusters):
        # i_func: cluster function index
        # f_prototype: i_func-th function on prototype cluster
        for i_func, f_prototype in enumerate(functions[i_orbit][0]):
            # Allow filtering by linear_function_index
            if (
                linear_function_indices is None
                or linear_function_index in linear_function_indices
            ):
                # prototype_functions: i_func-th function on prototype cluster only
                # orbit_functions: i_func-th function on each cluster in orbit
                prototype_functions = []
                orbit_functions = []
                for i_equiv, functions_by_cluster in enumerate(functions[i_orbit]):
                    if i_equiv == 0:
                        prototype_functions.append(functions_by_cluster[i_func])
                    orbit_functions.append(functions_by_cluster[i_func])

                _add_variables_needed(
                    data=variables_needed,
                    functions=orbit_functions,
                    prim_neighbor_list=prim_neighbor_list,
                    occ_site_functions=occ_site_functions,
                )

                orbit_bfuncs.append(
                    # make global correlation contribution formulas
                    {
                        "linear_function_index": linear_function_index,
                        "linear_orbit_index": linear_orbit_index,
                        "cpp": orbit_bfunc_cpp_str(
                            orbit_functions=orbit_functions,
                            orbit_size=len(orbit_functions),
                            prim_neighbor_list=prim_neighbor_list,
                            cpp_fmt=cpp_fmt,
                        ),
                        "latex_prototype": orbit_bfunc_cpp_str(
                            # only functions from prototype cluster
                            orbit_functions=prototype_functions,
                            # orbit_size=len(orbit_functions) -> same normalization
                            orbit_size=len(orbit_functions),
                            prim_neighbor_list=prim_neighbor_list,
                            cpp_fmt=cpp_fmt,
                            mode="latex",
                            label_site_using="cluster_site_index",
                        ),
                        "latex_prototype_with_neighbor_indices": orbit_bfunc_cpp_str(
                            # only functions from prototype cluster
                            orbit_functions=prototype_functions,
                            # orbit_size=len(orbit_functions) -> same normalization
                            orbit_size=len(orbit_functions),
                            prim_neighbor_list=prim_neighbor_list,
                            cpp_fmt=cpp_fmt,
                            mode="latex",
                            label_site_using="neighborhood_site_index",
                        ),
                        "latex_orbit": orbit_bfunc_cpp_str(
                            orbit_functions=orbit_functions,
                            orbit_size=len(orbit_functions),
                            prim_neighbor_list=prim_neighbor_list,
                            cpp_fmt=cpp_fmt,
                            mode="latex",
                            label_site_using="neighborhood_site_index",
                        ),
                    }
                )
            linear_function_index += 1
        linear_orbit_index += 1
    return (orbit_bfuncs, variables_needed)


def _print_latex_orbit_bfuncs(builder):
    """For debugging purposes, print the orbit basis functions in LaTeX format."""
    orbit_bfuncs, variables_needed = make_orbit_bfuncs(
        prim_neighbor_list=builder.prim_neighbor_list,
        clusters=builder.clusters,
        functions=builder.functions,
        occ_site_functions=builder.occ_site_functions,
        cpp_fmt=CppFormatProperties(),
        linear_function_indices=None,
    )

    print(
        "occ_var_name:",
        builder.occ_site_functions_info.get("occ_var_name").format(b="b", m="m"),
    )
    print("occ_var_desc:", builder.occ_site_functions_info.get("occ_var_desc"))

    for orbit_bfunc in orbit_bfuncs:
        linear_function_index = orbit_bfunc.get("linear_function_index")
        orbit_latex_formula = orbit_bfunc.get("latex_orbit")
        prototype_latex_formula = orbit_bfunc.get("latex_prototype")
        print("Linear function index: ", linear_function_index)
        print("prototype_formula: ", prototype_latex_formula)
        print("formula: ", orbit_latex_formula)
        print()


def _make_site_bfuncs_data(
    linear_function_index: int,
    linear_orbit_index: int,
    equiv_functions: list[list[PolynomialFunction]],
    orbit_size: int,
    prim_neighbor_list: PrimNeighborList,
    cpp_fmt: CppFormatProperties,
) -> dict:
    data = {
        "linear_function_index": linear_function_index,
        "linear_orbit_index": linear_orbit_index,
        "at": [],
    }
    for neighbor_list_index, equiv_functions_by_point in enumerate(equiv_functions):
        data["at"].append(
            {
                "neighbor_list_index": neighbor_list_index,
                "n_point_functions": len(equiv_functions_by_point),
                "cpp": site_bfunc_cpp_str(
                    point_functions=equiv_functions_by_point,
                    orbit_size=orbit_size,
                    prim_neighbor_list=prim_neighbor_list,
                    cpp_fmt=cpp_fmt,
                ),
                "latex": site_bfunc_cpp_str(
                    point_functions=equiv_functions_by_point,
                    orbit_size=orbit_size,
                    prim_neighbor_list=prim_neighbor_list,
                    cpp_fmt=cpp_fmt,
                    mode="latex",
                ),
                "occ_delta_cpp": occ_delta_site_bfunc_cpp_str(
                    neighbor_list_index=neighbor_list_index,
                    point_functions=equiv_functions_by_point,
                    orbit_size=orbit_size,
                    prim_neighbor_list=prim_neighbor_list,
                    cpp_fmt=cpp_fmt,
                ),
                "occ_delta_latex": occ_delta_site_bfunc_cpp_str(
                    neighbor_list_index=neighbor_list_index,
                    point_functions=equiv_functions_by_point,
                    orbit_size=orbit_size,
                    prim_neighbor_list=prim_neighbor_list,
                    cpp_fmt=cpp_fmt,
                    mode="latex",
                ),
            }
        )
    return data


def make_site_bfuncs(
    prim: casmconfig.Prim,
    is_periodic: bool,
    prim_neighbor_list: PrimNeighborList,
    clusters: list[list[casmclust.Cluster]],
    functions: list[list[list[PolynomialFunction]]],
    occ_site_functions: list[dict],
    cpp_fmt: CppFormatProperties,
    linear_function_indices: Optional[set[int]] = None,
    make_variable_name_f: Optional[Callable] = None,
    local_discrete_dof: Optional[dict] = None,
) -> tuple[list[dict], list[dict[str, list[list[int]]]]]:
    """Convert clusters and functions to data used by Jinja templates to write the \
    point correlation evaluating methods.

    Notes
    -----

    - For local basis sets, this currently returns empty lists

    Parameters
    ----------
    prim: libcasm.configuration.Prim
        The prim.
    is_periodic: bool
        Whether the Clexulator is for a periodic basis set or local basis set.
    prim_neighbor_list: libcasm.clexulator.PrimNeighborList
        :class:`~libcasm.clexulator.PrimNeighborList`, containing all the neighbors
        needed for evaluation, as obtained from
        :class:`~casm.bset.cluster_functions.ClusterFunctionBuilder`.
    clusters: list[list[libcasm.clusterography.Cluster]]
        List of clusters, as generated by
        :class:`~casm.bset.cluster_functions.ClusterFunctionBuilder`. Site order
        must be consistent with `functions`.
    functions: list[list[list[casm.bset.polynomial_functions.PolynomialFunction]]]
        List of functions, as generated by
        :class:`~casm.bset.cluster_functions.ClusterFunctionBuilder`.
    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 = functions[function_index][occupant_index]``.

    cpp_fmt: CppFormatProperties
        C++ string formatting properties. If None, the default values are used.
    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
    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).

    Returns
    -------
    site_bfuncs_data: 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

    variables_needed_at: list[dict[str,list[list[int]]]]
        For each point, the variables needed to evaluate the point functions. The value
        of `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.
    """

    # set n_point_corr_sites_max
    if is_periodic:
        n_point_corr_sites_max = len(prim_neighbor_list.sublattice_indices())
    else:
        n_point_corr_sites_max = prim_neighbor_list.n_neighborhood_sites()

    ### Start: constant term, linear_function_index==0 ###
    linear_function_index = 0
    linear_orbit_index = 0
    site_bfuncs = []
    variables_needed_at = [{} for i in range(n_point_corr_sites_max)]

    ### For constant term, linear_function_index==0:
    # Allow filtering by linear_function_index
    if (
        linear_function_indices is None
        or linear_function_index in linear_function_indices
    ):
        site_bfuncs.append(
            {
                "linear_function_index": linear_function_index,
                "linear_orbit_index": linear_orbit_index,
                "at": [],
            }
        )

    linear_function_index += 1

    ### For all cluster orbits (includes null cluster):
    # i_orbit: cluster orbit index
    # orbit: list[Cluster]
    for i_orbit, orbit in enumerate(clusters):
        # point_functions: list[list[list[PolynomialFunction]]]
        #     All point functions for this cluster orbit, for each point corr site
        #     point_functions[i_func][nlist_index][i_point_function]
        if is_periodic:
            # take global functions (non-duplicating) and make all point functions
            point_functions = make_point_functions(
                prim=prim,
                prim_neighbor_list=prim_neighbor_list,
                orbit=orbit,
                orbit_functions=functions[i_orbit],
                make_variable_name_f=make_variable_name_f,
                local_discrete_dof=local_discrete_dof,
            )
        else:
            point_functions = make_local_point_functions(
                prim_neighbor_list=prim_neighbor_list,
                orbit=orbit,
                orbit_functions=functions[i_orbit],
            )

        # i_func: cluster function index
        # equiv_functions: list[list[PolynomialFunction]]
        #     All equivalent point functions, for each point corr site
        #     equiv_functions[nlist_index][i_point_function]
        for i_func, equiv_functions in enumerate(point_functions):
            # Allow filtering by linear_function_index
            if (
                linear_function_indices is None
                or linear_function_index in linear_function_indices
            ):
                for neighbor_list_index, equiv_functions_by_point in enumerate(
                    equiv_functions
                ):
                    _add_variables_needed(
                        data=variables_needed_at[neighbor_list_index],
                        functions=equiv_functions_by_point,
                        prim_neighbor_list=prim_neighbor_list,
                        occ_site_functions=occ_site_functions,
                    )

                site_bfuncs.append(
                    # make point correlation formulas for each point corr site
                    _make_site_bfuncs_data(
                        linear_function_index=linear_function_index,
                        linear_orbit_index=linear_orbit_index,
                        equiv_functions=equiv_functions,
                        orbit_size=len(orbit),
                        prim_neighbor_list=prim_neighbor_list,
                        cpp_fmt=cpp_fmt,
                    )
                )
            linear_function_index += 1
        linear_orbit_index += 1
    return (site_bfuncs, variables_needed_at)


def _print_latex_site_bfuncs(builder):
    """For debugging purposes, print the site basis functions in LaTeX format."""

    is_periodic = builder._phenomenal is not None
    site_bfuncs, site_bfuncs_variables_needed_at = make_site_bfuncs(
        prim=builder._prim,
        is_periodic=is_periodic,
        prim_neighbor_list=builder.prim_neighbor_list,
        clusters=builder.clusters,
        functions=builder.functions,
        occ_site_functions=builder.occ_site_functions,
        cpp_fmt=CppFormatProperties(),
        linear_function_indices=None,
        make_variable_name_f=builder._make_variable_name_f,
        local_discrete_dof=builder.local_discrete_dof,
    )

    print(
        "occ_var_name:",
        builder.occ_site_functions_info.get("occ_var_name").format(b="b", m="m"),
    )
    print("occ_var_desc:", builder.occ_site_functions_info.get("occ_var_desc"))

    for site_bfunc_all in site_bfuncs:
        linear_function_index = site_bfunc_all.get("linear_function_index")
        print("### Linear function index: ", linear_function_index, "###")
        print()

        for site_bfunc_at in site_bfunc_all.get("at"):
            neighbor_list_index = site_bfunc_at.get("neighbor_list_index")
            latex_formula = site_bfunc_at.get("latex")

            print(
                "Linear function index: ",
                linear_function_index,
                "at: ",
                neighbor_list_index,
            )
            print("formula: ", latex_formula)
            occ_delta_latex_formula = site_bfunc_at.get("occ_delta_latex")
            print("occ_delta_formula: ", occ_delta_latex_formula)
            print()


class ClexulatorWriter:
    """Write Clexulator source files and related files

    The design is:

    - :func:`~casm.bset.write_clexulator` is the intended user interface for
      writing Clexulator source files
    - :class:`~casm.bset.clexwriter.ClexulatorWriter` implements the things that
      all Clexulator writer versions need to do (write the C++ source files,
      basis.json, equivalents_info.json, etc.)
    - :class:`~casm.bset.clexwriter._methods.WriterV1Basic` or other classes
      (WriterV1Diff, etc.) implements the specifics for a particular Clexulator
      `version`, by specifying the template files and generating the variables used
      by the templates.

    """

    def __init__(
        self,
        bset_dir: pathlib.Path,
        version: str,
        project_name: str,
        bset_name: str,
        linear_function_indices: Optional[set[int]] = None,
        cpp_fmt: Optional[CppFormatProperties] = None,
    ):
        """

        .. rubric:: Constructor

        Parameters
        ----------
        bset_dir: pathlib.Path
            The path to the basis set directory where the Clexulator and related files
            should be written

        version: str
            The Clexulator version to write. One of:

            - "v1.basic": Standard CASM v1 compatible Clexulator, without automatic
              differentiation
            - "v1.diff": (TODO) CASM v1 compatible Clexulator, with ``fadbad`` automatic
              differentiation enabled

        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.

        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.
            Generally this is not known the first time a Clexulator is generated, but
            after fitting coefficients it may be used to re-generate the Clexulator
            with the subset of the basis functions needed.

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

        """
        self.bset_dir = pathlib.Path(bset_dir)
        """pathlib.Path: The path to the basis set directory where the Clexulator and \
        related files should be written."""

        self.version = version
        """str: The Clexulator version to write. 
        
        One of:
        
        - "v1.basic": Standard CASM v1 compatible Clexulator, without automatic
            differentiation.
        - "v1.diff": (TODO) CASM v1 compatible Clexulator, with ``fadbad`` automatic
            differentiation enabled.
        
        """

        self.project_name = project_name
        """str: Project name."""

        self.bset_name = bset_name
        """str: Basis set name. 
        
        Used to construct the Clexulator class name. This must consist of alphanumeric 
        characters and underscores only."""

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

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

        self.src_path = None
        """Optional[pathlib.Path]: The path to the Clexulator source file, once \
        written."""

        self.local_src_path = None
        """Optional[list[pathlib.Path]]: The paths to the local Clexulator source \
        files, once written."""

        self.generated_files = None
        """Optional[list[pathlib.Path]]: The paths to the generated files, once \
        written."""

        if not re.match(
            R"^[a-zA-Z_]+\w*",
            project_name,
        ):
            raise Exception(
                f"project_name='{project_name}' is not valid: ",
                "Must consist alphanumeric characters and underscores only. "
                "The first character may not be a number.",
            )
        if not re.match(
            R"^\w+",
            bset_name,
        ):
            raise Exception(
                f"bset_name='{bset_name}' is not valid: ",
                "Must consist alphanumeric characters and underscores only.",
            )

        if self.version == "v1.basic":
            from ._write_v1_basic import WriterV1Basic

            self.writer_type = WriterV1Basic
        else:
            raise Exception(f"Error in ClexulatorWriter: Unknown version '{version}'")

    def write(
        self,
        prim: casmconfig.Prim,
        clex_basis_specs: ClexBasisSpecs,
        prim_neighbor_list: PrimNeighborList,
        verbose: bool = True,
        very_verbose: bool = False,
    ):
        """Write Clexulator source files and related files

        Parameters
        ----------
        prim: libcasm.configuration.Prim
            The prim, with symmetry information.

        clex_basis_specs: casm.bset.cluster_functions.ClexBasisSpecs
            Parameters specifying the cluster orbits and basis function type and order.

        prim_neighbor_list: PrimNeighborList
            The :class:`PrimNeighborList` is used to uniquely index sites with local
            variables included in the cluster functions, relative to a reference unit
            cell. If None, a default neighbor list is constructed.

        verbose: bool = True
            Print progress statements

        very_verbose: bool = False
            Print detailed progress statements from the cluster functions builder.

        """
        cluster_specs = clex_basis_specs.cluster_specs
        bfunc_specs = clex_basis_specs.basis_function_specs

        clusters = [orbit[0] for orbit in cluster_specs.make_orbits()]

        orbit_branch_max_poly_order = {
            int(key): value
            for key, value in bfunc_specs.orbit_branch_max_poly_order.items()
        }

        phenomenal = cluster_specs.phenomenal()
        is_periodic = phenomenal is None

        occ_site_basis_functions_specs = None
        if "occ" in bfunc_specs.dof_specs:
            if "site_basis_functions" in bfunc_specs.dof_specs["occ"]:
                occ_site_basis_functions_specs = bfunc_specs.dof_specs["occ"][
                    "site_basis_functions"
                ]

        # store a list of generated files
        generated_files = []

        if verbose:
            start = time.time()
            print("Building cluster functions...")
            sys.stdout.flush()

        builder = ClusterFunctionsBuilder(
            prim=prim,
            generating_group=cluster_specs.generating_group(),
            dofs=bfunc_specs.dofs,
            clusters=clusters,
            phenomenal=cluster_specs.phenomenal(),
            global_max_poly_order=bfunc_specs.global_max_poly_order,
            orbit_branch_max_poly_order=orbit_branch_max_poly_order,
            occ_site_basis_functions_specs=occ_site_basis_functions_specs,
            prim_neighbor_list=prim_neighbor_list,
            verbose=very_verbose,
        )

        path = self.bset_dir / "cluster_functions.json.gz"
        generated_files.append(path)
        dump(builder.to_dict(), path, quiet=True, force=True, gz=True)

        if verbose:
            print("Building cluster functions DONE")
            print(f"- n_orbits: {len(builder.clusters)}")
            print(f"- n_functions: {builder.n_functions}")
            elapsed_time = time.time() - start
            print(f"build time: {elapsed_time:0.4f} (s)")
            print()
            sys.stdout.flush()

        env = jinja2.Environment(
            trim_blocks=True,
            lstrip_blocks=True,
            loader=jinja2.PackageLoader("casm.bset"),
        )
        template = env.get_template(self.writer_type.template_name())

        ## common clexulator writer parameters
        writer_params = dict(
            project_name=self.project_name,
            bset_name=self.bset_name,
            is_periodic=is_periodic,
            prim_neighbor_list=builder.prim_neighbor_list,
            occ_site_functions=builder.occ_site_functions,
            occ_site_functions_info=builder.occ_site_functions_info,
            linear_function_indices=self.linear_function_indices,
            make_variable_name_f=builder._make_variable_name_f,
            local_discrete_dof=builder.local_discrete_dof,
            cpp_fmt=self.cpp_fmt,
            verbose=verbose,
        )

        ## write periodic clexulator / prototype local clexulator
        if verbose:
            start = time.time()
            print("Generating variables...")
            sys.stdout.flush()

        writer = self.writer_type(
            prim=prim,
            i_clex=None,
            clusters=builder.clusters,
            functions=builder.functions,
            **writer_params,
        )

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

        # Clexulator source file
        path = self.bset_dir / f"{writer.clexulator_name}.cc"
        self.src_path = path
        generated_files.append(path)
        with open(path, "w") as f:
            f.write(template.render(writer.variables()))

        # basis.json file
        path = self.bset_dir / "basis.json"
        generated_files.append(path)
        data = builder.basis_dict(
            clex_basis_specs=clex_basis_specs, coordinate_mode="frac"
        )
        dump(data, path, quiet=True, force=True)

        # write all template variables
        path = self.bset_dir / "variables.json.gz"
        generated_files.append(path)
        dump(writer.variables(), path, quiet=True, force=True, gz=True)

        ## if local clexulators
        if phenomenal is not None:
            # write equivalents_info.json
            path = self.bset_dir / "equivalents_info.json"
            generated_files.append(path)
            dump(builder.equivalents_info_dict(), path, quiet=True, force=True)

            # write each local clexulator
            self.local_src_path = []
            for i_clex in range(len(builder.equivalent_clusters)):
                if verbose:
                    start = time.time()
                    print(f"Generating variables for local clexulator {i_clex}...")
                    sys.stdout.flush()

                writer = self.writer_type(
                    prim=prim,
                    i_clex=i_clex,
                    clusters=builder.equivalent_clusters[i_clex],
                    functions=builder.equivalent_functions[i_clex],
                    **writer_params,
                )

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

                local_dir = self.bset_dir / f"{i_clex}"
                local_dir.mkdir(parents=True, exist_ok=True)

                # write local Clexulator
                path = local_dir / f"{writer.clexulator_name}.cc"
                self.local_src_path.append(path)
                generated_files.append(path)
                with open(self.local_src_path[-1], "w") as f:
                    f.write(template.render(writer.variables()))

                # write all template variables
                path = local_dir / "variables.json"
                generated_files.append(path)
                dump(writer.variables(), path, quiet=True, force=True)

        # write generated_files.json (includes itself)
        path = self.bset_dir / "generated_files.json"
        generated_files.append(path)
        self.generated_files = generated_files
        data = {
            "src_path": str(self.src_path.relative_to(self.bset_dir)),
            "all": [str(p.relative_to(self.bset_dir)) for p in self.generated_files],
        }
        if self.local_src_path is not None:
            data["local_src_path"] = [
                str(p.relative_to(self.bset_dir)) for p in self.local_src_path
            ]
        dump(data, path, quiet=True, force=True)
        return None
