import libcasm.clusterography as casmclust
import libcasm.configuration as casmconfig
import libcasm.occ_events._occ_events as _occ_events
import libcasm.sym_info as sym_info
import libcasm.xtal as xtal


def orbits_to_dict(
    orbits: list[list[casmclust.Cluster]],
    prim: casmconfig.Prim,
) -> dict:
    """Represent a list of cluster orbits as a list of Python dict

    Parameters
    ----------
    orbits: list[list[casmclust.Cluster]]
        List of cluster orbits. Orbits may not be empty.
    prim: libcasm.configuration.Prim
        The prim, used to generate site distances.

    Returns
    -------
    data: list[dict]
        A list of Python dict, one per orbit, in the
        `Cluster Orbit format <https://prisms-center.github.io/CASMcode_docs/formats/casm/clusterography/ClusterOrbits/#cluster-orbit-json-object>`_.

    """
    orbits_data = []
    for i_orbit, orbit in enumerate(orbits):
        orbit_data = {
            "linear_orbit_index": i_orbit,
            "mult": len(orbit),
            "prototype": orbit[0].to_dict(xtal_prim=prim.xtal_prim),
            "elements": [],
        }

        element_data = orbit_data["elements"]
        for cluster in orbit:
            element_data.append(cluster.to_dict(xtal_prim=prim.xtal_prim))

        orbits_data.append(orbit_data)
    return orbits_data


def make_symop_inverse(
    op: xtal.SymOp,
) -> xtal.SymOp:
    """Make the inverse SymOp

    Parameters
    ----------
    op: xtal.SymOp
        A SymOp

    Returns
    -------
    op_inv: xtal.SymOp
        The inverse SymOp
    """
    return xtal.SymOp(
        matrix=op.matrix().T,
        translation=-(op.matrix().T @ op.translation()),
        time_reversal=op.time_reversal(),
    )


def make_equivalents_generators(
    phenomenal: casmclust.Cluster,
    generating_group: sym_info.SymGroup,
    prim: casmconfig.Prim,
) -> tuple[list[xtal.SymOp], list[int], list[xtal.IntegralSiteCoordinateRep]]:
    """Make symmetry operations that generate all the equivalent local orbits in \
    the primitive cell.

    Notes
    -----

    - If the `cluster_specs.generating_group()` is a subgroup of the phenomenal
      cluster group, then there will be >1 distinct sets of local orbits around the
      phenomenal cluster.
    - This method finds the cluster group operations that generate the distinct set of
      local orbits, then it uses the phenomenal cluster orbit equivalence map to
      find operations that transform the local orbits sets to the equivalent
      phenomenal clusters in the primitive cell.

    Parameters
    ----------
    phenomenal: casmclust.Cluster
        The phenomenal cluster. The phenomenal cluster must be chosen from one of the
        equivalents that is generated by
        :func:`~libcasm.clusterography.make_periodic_orbit` using the prim factor
        group.
    generating_group: sym_info.SymGroup
        The local orbits generating group.
    prim: casmconfig.Prim
        The prim, with symmetry info

    Returns
    -------
    ops: list[libcasm.xtal.SymOp]
        Symmetry operations that generate the equivalent local orbits in
        the primitive unit cell.

    indices: list[int]
        Indices of prim factor group operations corresponding to `ops`.

    site_reps: list[libcasm.xtal.IntegralSiteCoordinateRep]
        Symmetry group representation for transforming IntegralSiteCoordinate
        and Cluster corresponding to `ops`.
    """

    # collect needed sym groups
    factor_grp = prim.factor_group
    generating_grp = generating_group
    site_rep = prim.integral_site_coordinate_symgroup_rep

    phenomenal_cluster_grp = casmclust.make_cluster_group(
        cluster=phenomenal,
        group=factor_grp,
        lattice=prim.xtal_prim.lattice(),
        integral_site_coordinate_symgroup_rep=site_rep,
    )

    # find indices of factor group operations
    # that generate equivalent (but rotated) local basis sets
    # about the phenomenal cluster.
    i_factor_grp_equiv_on_phenomenal = set()
    for i_cluster_grp in range(len(phenomenal_cluster_grp.elements)):
        i_factor_grp = phenomenal_cluster_grp.head_group_index[i_cluster_grp]

        i_factor_grp_min = i_factor_grp
        for j_generating_grp in range(len(generating_grp.elements)):
            j_factor_grp = generating_grp.head_group_index[j_generating_grp]

            k_product = factor_grp.mult(j_factor_grp, i_factor_grp)
            if k_product < i_factor_grp_min:
                i_factor_grp_min = k_product

        i_factor_grp_equiv_on_phenomenal.add(i_factor_grp_min)

    # generate the phenomenal cluster orbit and equivalence map
    phenomenal_orbit = casmclust.make_periodic_orbit(
        orbit_element=phenomenal,
        integral_site_coordinate_symgroup_rep=site_rep,
    )
    equivalence_map = casmclust.make_periodic_equivalence_map(
        orbit=phenomenal_orbit,
        symgroup=factor_grp,
        lattice=prim.xtal_prim.lattice(),
        integral_site_coordinate_symgroup_rep=site_rep,
    )
    equivalence_map_indices = casmclust.make_periodic_equivalence_map_indices(
        orbit=phenomenal_orbit,
        integral_site_coordinate_symgroup_rep=site_rep,
    )

    # find op and factor group index that transform phenomenal to orbit prototype
    to_prototype_op = None
    i_to_prototype = None
    sorted_phenomenal = phenomenal.sorted()
    for i_clust, clust in enumerate(phenomenal_orbit):
        sorted_clust = clust.sorted()
        if sorted_clust == sorted_phenomenal:
            to_prototype_op = make_symop_inverse(equivalence_map[i_clust][0])
            i_to_prototype = factor_grp.inv(equivalence_map_indices[i_clust][0])
            break
    if to_prototype_op is None:
        raise Exception(
            "Error in make_equivalents_generators: Failed to find to_prototype_op."
            "The phenomenal cluster must be chosen from one of the"
            "equivalents generated by `libcasm.clusterography.make_periodic_orbit`."
        )

    # equiv bset = to_equiv_op * (to_prototype_op * (cluster_grp_op * phenomenal bset))
    #                                              ^ distinct bset on phenomenal
    #                             ^ bset on prototype
    #              ^ bset on other phenomenal
    generating_indices = []
    for i_factor_grp in i_factor_grp_equiv_on_phenomenal:
        for i_equiv in range(len(equivalence_map)):
            i_to_equiv = equivalence_map_indices[i_equiv][0]
            i_generating_op = factor_grp.mult(
                i_to_equiv, factor_grp.mult(i_to_prototype, i_factor_grp)
            )

            generating_indices.append(i_generating_op)
    generating_indices.sort()

    generating_ops = []
    for i in generating_indices:
        generating_ops.append(factor_grp.elements[i])

    generating_site_reps = casmclust.make_integral_site_coordinate_symgroup_rep(
        group_elements=generating_ops,
        xtal_prim=prim.xtal_prim,
    )

    return (generating_ops, generating_indices, generating_site_reps)


def make_neighborhood(
    clusters: list[list[casmclust.Cluster]],
):
    """Return a list of unique sites from orbits of clusters

    Parameters
    ----------
    clusters: list[list[libcasm.clusterography.Cluster]]
        Clusters, arranged in orbits.

    Returns
    -------
    neighborhood: list[libcasm.xtal.IntegralSiteCoordinate]
        The neighborhood of sites in clusters. Excludes
    """
    neighborhood = []
    for i_orbit, orbit in enumerate(clusters):
        for equiv in orbit:
            for site in equiv:
                if site not in neighborhood:
                    neighborhood.append(site)
    return neighborhood


def make_occevent_cluster_specs(
    prim: casmconfig.Prim,
    phenomenal_occ_event: _occ_events.OccEvent,
    max_length: list[float],
    cutoff_radius: list[float],
    custom_generators: list[casmclust.ClusterOrbitGenerator] = [],
) -> casmclust.ClusterSpecs:
    """Construct ClusterSpecs for local-cluster orbits around an OccEvent

    Parameters
    ----------
    prim: libcasm.configuration.Prim
        The prim, with symmetry information
    phenomenal_occ_event: ~libcasm.occ_events.OccEvent
        The orbit generating group is the subgroup of the prim factor
        group that leaves `phenomenal_occ_event` invariant.
    max_length: list[float]
        The maximum site-to-site distance to allow in clusters, by number
        of sites in the cluster. Example: `[0.0, 0.0, 5.0, 4.0]` specifies
        that pair clusters up to distance 5.0 and triplet clusters up to
        distance 4.0 should be included. The null cluster and point
        cluster values (elements 0 and 1) are arbitrary.
    cutoff_radius: list[float]
        For local clusters, the maximum distance of sites from any
        phenomenal cluster site to include in the local environment, by
        number of sites in the cluster. The null cluster value
        (element 0) is arbitrary.
    custom_generators: list[~libcasm.clusterography.ClusterOrbitGenerator]=[]
          Specifies clusters that should be uses to construct orbits
          regardless of the max_length or cutoff_radius parameters

    Returns
    -------
    cluster_specs: ~libcasm.clusterography.ClusterSpecs
        The resulting ClusterSpecs
    """
    symgroup_rep = _occ_events.make_occevent_symgroup_rep(
        prim.factor_group.elements, prim.xtal_prim
    )
    occevent_group = _occ_events.make_occevent_group(
        occ_event=phenomenal_occ_event,
        group=prim.factor_group,
        lattice=prim.xtal_prim.lattice(),
        occevent_symgroup_rep=symgroup_rep,
    )
    return casmclust.ClusterSpecs(
        xtal_prim=prim.xtal_prim,
        generating_group=occevent_group,
        max_length=max_length,
        phenomenal=phenomenal_occ_event.cluster(),
        cutoff_radius=cutoff_radius,
        custom_generators=custom_generators,
    )
