# General imports.
import hashlib

# Digichem imports.
from digichem.result.result import Result_set
from digichem.result.metadata import Metadata
from digichem.result.atom import Atom_list
from digichem.result.ground_state import Ground_state
from digichem.result.excited_state import Excited_state_list
from digichem.result.emission import Relaxed_excited_state
from digichem.result.orbital import Molecular_orbital_list
from digichem.result.alignment.base import Alignment, Minimal


class Merged(Result_set):
    """
    A type of result set that represents one or more separate calculations merged into one result set.
    """
    
    def __init__(self, results, *args, vertical_emission = None, adiabatic_emission = None, **kwargs):
        """
        Constructor for Merged result sets.
        :param vertical_emission: An optional dictionary of Relaxed_excited_state objects representing vertical emission energies (one for each multiplicity).
        :param adiabatic_emission: An optional dictionary of Relaxed_excited_state objects representing vertical adiabatic energies (one for each multiplicity).
        """
        super().__init__(*args, **kwargs)
        self.results = results
        self.emission.vertical = vertical_emission if vertical_emission is not None else {}
        self.emission.adiabatic = adiabatic_emission if adiabatic_emission is not None else {}
            
    @classmethod
    def from_results(self, *results, options):
        """
        Create a Merged result set object from a number of result sets.
        
        :param *results: List of result sets to merge.
        :param options: A Digichem config nested dictionary containing options to control parsing.
        """
        # First, get a merged metadata object.
        metadatas = [result.metadata for result in results]
        merged_metadata = Metadata.merge(*metadatas)

        # 'Merge' our atoms.
        raw_atoms = Atom_list.merge(*[result.raw_atoms for result in results], charge = merged_metadata.charge)
        # And alignment.
        alignment_class = Alignment.from_class_handle(options['alignment']) if options['alignment'] is not None else Minimal
        atoms = alignment_class(raw_atoms, charge = merged_metadata.charge)
        
        # Use special method for merging orbitals.
        orbitals, beta_orbitals = Molecular_orbital_list.merge_orbitals([result.orbitals for result in results], [result.beta_orbitals for result in results])
        
        # Merge NMR (which depends on atoms and options).
        nmr = type(getattr(results[0], "nmr")).merge(*[getattr(result, "nmr") for result in results], atoms = atoms, options = options)
        
        # Merge remaining attributes.
        attrs = {}
        for attr in ["energies", "pdm", "excited_states", "vibrations", "soc"]:
            attrs[attr] = type(getattr(results[0], attr)).merge(*[getattr(result, attr) for result in results])
        
        # Get a new ground state.
        ground_state = Ground_state.from_energies(merged_metadata.charge, merged_metadata.multiplicity, attrs['energies'])
        
        # Build new list of energy states.
        energy_states = Excited_state_list()
        energy_states.append(ground_state)
        energy_states.extend(attrs['excited_states'])
        
        # Get a new ID for the merged object.
        # This is generated by combining the hashes of each of the contributing results.
        hasher = hashlib.sha1()
        
        for result in results:
            if result._id:
                byte_hash = bytes.fromhex(result._id)
                
            hasher.update(byte_hash)
            
        _id = hasher.hexdigest()
        
        merged_results =  self(
            _id = _id,
            metadata = merged_metadata,
            results = results,
            ground_state = ground_state,
            raw_atoms = raw_atoms,
            atoms = atoms,
            energy_states = energy_states,
            orbitals = orbitals,
            beta_orbitals = beta_orbitals,
            nmr = nmr,
            **attrs
            )
        
        # Try and guess emission.
        merged_results.emission.vertical, merged_results.emission.adiabatic = Relaxed_excited_state.guess_from_results(*results)
        
        # Done.
        return merged_results
