import numpy as np


class OptimizationStep:
    """Each instantiation of this class corresponds to a step in the optimization process.

    Attributes
    ----------
    structure: riigid.Structure
        The initial structure of the optimization step
    force_on_atoms: numpy.ndarray of shape (n_atoms_in_structure, 3)
        Forces acting on structure.atoms; [eV/Å]
    energy: number
        The energy of structure; [eV]
    update_structure: riigid.Structure
        The updated structure (generated by letting the forces act on structure)
    force_on_fragments: list of numpy.ndarrays of shape (3,)
        Net force on each fragment in self.structure; [eV/Å]
    torque_on_fragments: list of numpy.ndarrays of shape (3,)
        Torque on each fragment in self.structure; [eV]
    max_force_on_atom: (float, int)
        - Maximal force on an atom in self.structure; [eV/Å]
        - Index of the corresponding atom.

    max_force_on_fragment: (float, int)
        - Maximal force on a fragment in self.structure; [eV/Å]
        - Index of the corresponding fragment.

    max_torque_on_fragment: (float, int)
        - Maximal torque on a fragment in self.structure; [eV]
        - Index of the corresponding fragment.


    Note
    ----
    The stored forces and the energy belong to "structure", not to "updated_structure".

    """

    def __init__(self, structure, force_on_atoms, energy, updated_structure=None):
        """Initialize the optimization step.

        Parameters
        ----------
        structure: riigid.Structure
            The initial structure of the optimization step
        force_on_atoms: numpy.ndarray of shape (n_atoms_in_structure, 3)
            Forces acting on structure.atoms; [eV/Å]
        energy: number
            The energy of structure; [eV]
        update_structure: riigid.Structure
            The updated structure (generated by letting the forces act on structure)

        """
        self.structure = structure
        self.force_on_atoms = force_on_atoms
        self.energy = energy
        self.updated_structure = updated_structure

        # Calculate force and torque on fragments
        self.add_force_and_torque_on_fragments()

        # Calc max force/torque
        self.calc_max_force_and_torque()

    def add_updated_structure(self, updated_structure):
        """Add/change the updated structure of the optimization step.

        Parameters
        ----------
        update_structure: riigid.Structure
            The updated structure (generated by letting the forces act on structure)

        """
        self.updated_structure = updated_structure

    def remove_updated_structure(self):
        """Remove the updated structure from the optimization step."""
        self.updated_structure = None

    def add_force_and_torque_on_fragments(self):
        """Calculate force and torque on fragments of self.structure."""
        self.force_on_fragments = self.structure.calculate_net_force_on_fragments(
            forces=self.force_on_atoms
        )
        self.torque_on_fragments = self.structure.calculate_torque_on_fragments(
            forces=self.force_on_atoms
        )

    def calc_max_force_and_torque(self):
        """Find max (and argmax) force on atoms and fragments and torque on fragments."""
        # Calculate max force on an atom
        tmp = np.sqrt(np.sum((self.force_on_atoms) ** 2, axis=1))
        self.max_force_on_atom = (np.max(tmp), np.argmax(tmp))

        # Calculate max force on a fragment
        tmp = [np.sqrt(np.sum(tmp2**2)) for tmp2 in self.force_on_fragments]
        self.max_force_on_fragment = (np.max(tmp), np.argmax(tmp))

        # Calculate max torque on a fragment
        tmp = [np.sqrt(np.sum(tmp2**2)) for tmp2 in self.torque_on_fragments]
        self.max_torque_on_fragment = (np.max(tmp), np.argmax(tmp))
