# fmt: off

from ase.calculators.calculator import (
    BaseCalculator,
    CalculatorSetupError,
    all_changes,
)
from ase.stress import full_3x3_to_voigt_6_stress


class Mixer:
    def __init__(self, calcs, weights):
        self.check_input(calcs, weights)
        common_properties = set.intersection(
            *(set(calc.implemented_properties) for calc in calcs)
        )
        self.implemented_properties = list(common_properties)
        self.calcs = calcs
        self.weights = weights

    @staticmethod
    def check_input(calcs, weights):
        if len(calcs) == 0:
            raise CalculatorSetupError("Please provide a list of Calculators")
        if len(weights) != len(calcs):
            raise ValueError(
                "The length of the weights must be the same as"
                " the number of Calculators!"
            )

    def get_properties(self, properties, atoms):
        results = {}

        def get_property(prop):
            contribs = [calc.get_property(prop, atoms) for calc in self.calcs]
            # ensure that the contribution shapes are the same for stress prop
            if prop == "stress":
                shapes = [contrib.shape for contrib in contribs]
                if not all(shape == shapes[0] for shape in shapes):
                    if prop == "stress":
                        contribs = self.make_stress_voigt(contribs)
                    else:
                        raise ValueError(
                            f"The shapes of the property {prop}"
                            " are not the same from all"
                            " calculators"
                        )
            results[f"{prop}_contributions"] = contribs
            results[prop] = sum(
                weight * value for weight, value in zip(self.weights, contribs)
            )

        for prop in properties:  # get requested properties
            get_property(prop)
        for prop in self.implemented_properties:  # cache all available props
            if all(prop in calc.results for calc in self.calcs):
                get_property(prop)
        return results

    @staticmethod
    def make_stress_voigt(stresses):
        new_contribs = []
        for contrib in stresses:
            if contrib.shape == (6,):
                new_contribs.append(contrib)
            elif contrib.shape == (3, 3):
                new_cont = full_3x3_to_voigt_6_stress(contrib)
                new_contribs.append(new_cont)
            else:
                raise ValueError(
                    "The shapes of the stress"
                    " property are not the same"
                    " from all calculators"
                )
        return new_contribs


class LinearCombinationCalculator(BaseCalculator):
    """Weighted summation of multiple calculators."""

    def __init__(self, calcs, weights):
        """Implementation of sum of calculators.

        calcs: list
            List of an arbitrary number of :mod:`ase.calculators` objects.
        weights: list of float
            Weights for each calculator in the list.
        """
        super().__init__()
        self.mixer = Mixer(calcs, weights)
        self.implemented_properties = self.mixer.implemented_properties

    def calculate(self, atoms, properties, system_changes):
        """Calculates all the specific property for each calculator and
        returns with the summed value.

        """
        self.atoms = atoms.copy()  # for caching of results
        self.results = self.mixer.get_properties(properties, atoms)

    def __str__(self):
        calculators = ", ".join(
            calc.__class__.__name__ for calc in self.mixer.calcs
        )
        return f"{self.__class__.__name__}({calculators})"


class MixedCalculator(LinearCombinationCalculator):
    """
    Mixing of two calculators with different weights

    H = weight1 * H1 + weight2 * H2

    Has functionality to get the energy contributions from each calculator

    Parameters
    ----------
    calc1 : ASE-calculator
    calc2 : ASE-calculator
    weight1 : float
        weight for calculator 1
    weight2 : float
        weight for calculator 2
    """

    def __init__(self, calc1, calc2, weight1, weight2):
        super().__init__([calc1, calc2], [weight1, weight2])

    def set_weights(self, w1, w2):
        self.mixer.weights[0] = w1
        self.mixer.weights[1] = w2

    def get_energy_contributions(self, atoms=None):
        """Return the potential energy from calc1 and calc2 respectively"""
        self.calculate(
            properties=["energy"],
            atoms=atoms,
            system_changes=all_changes
        )
        return self.results["energy_contributions"]


class SumCalculator(LinearCombinationCalculator):
    """SumCalculator for combining multiple calculators.

    This calculator can be used when there are different calculators
    for the different chemical environment or for example during delta
    leaning. It works with a list of arbitrary calculators and
    evaluates them in sequence when it is required.  The supported
    properties are the intersection of the implemented properties in
    each calculator.

    """

    def __init__(self, calcs):
        """Implementation of sum of calculators.

        calcs: list
            List of an arbitrary number of :mod:`ase.calculators` objects.
        """

        weights = [1.0] * len(calcs)
        super().__init__(calcs, weights)


class AverageCalculator(LinearCombinationCalculator):
    """AverageCalculator for equal summation of multiple calculators (for
    thermodynamic purposes)."""

    def __init__(self, calcs):
        """Implementation of average of calculators.

        calcs: list
            List of an arbitrary number of :mod:`ase.calculators` objects.
        """
        n = len(calcs)

        if n == 0:
            raise CalculatorSetupError(
                "The value of the calcs must be a list of Calculators"
            )

        weights = [1 / n] * n
        super().__init__(calcs, weights)