import h5py
import numpy as np
from pyscf.symm.param import IRREP_ID_MOLPRO
from pyscf.symm.geom import detect_symm, get_subgroup
from pyscf.symm import label_orb_symm, irrep_name2id, irrep_id2name
from pyscf.lib import logger, tag_array
from pyscf import gto
from pyscf.tools import molden

CARTESIAN = "Cartesian or mixed Cartesian-spherical basis sets"

def get_mo_from_h5 (mol, h5fname, symmetry=None):
    ''' Get MO vectors for a pyscf molecule from an h5 file written by OpenMolcas

        Args:
            mol : instance gto.mole
                Must be in the same point group as the OpenMolcas calculation,
                or set the symmetry argument
            h5fname : str
                Path to an h5 file generated by OpenMolcas containing (at
                least) groups 'BASIS_FUNCTION_IDS', 'MO_OCCUPATIONS',
                'MO_VECTORS', and 'MO_ENERGIES' and additionally 'DESYM_MATRIX'
                and 'DESYM_BASIS_FUNCTION_IDS' if symmetry is used.

        Kwargs:
            symmetry : str
                Point group of the calculation in OpenMolcas. If not provided,
                mol.groupname is used instead

        Returns:
            mo_coeff : ndarray of shape (nao_nr, nao_nr)
    '''

    if symmetry is not None:
        mol = mol.copy()
        mol.build(symmetry=symmetry)
    nao = mol.nao_nr ()

    with h5py.File (h5fname, 'r') as f:
        try:
            molcas_basids = f['DESYM_BASIS_FUNCTION_IDS'][()]
            molcas_usymm = f['DESYM_MATRIX'][()].reshape (nao, nao)
        except KeyError:
            assert (not mol.symmetry), "Can't find desym_ data; mol.symmetry = {}".format (mol.symmetry)
            molcas_basids = f['BASIS_FUNCTION_IDS'][()]
        molcas_coeff = f['MO_VECTORS'][()]
        mo_energy = f['MO_ENERGIES'][()]
        mo_occ = f['MO_OCCUPATIONS'][()]

    idx_ao = []
    for (c, n, l, m) in molcas_basids:
        if l < 0: raise NotImplementedError (CARTESIAN)
        # 0-index atom list in PySCF, 1-index atom list in Molcas
        c -= 1
        # Actual principal quantum number in PySCF, 1-indexed list in Molcas
        n += l
        # l=1, ml=(-1,0,1) is (x,y,z) in PySCF, (y,z,x) in Molcas
        if l == 1:
            m = m - 2 if m > 0 else m + 1
        idx_ao.append (mol.search_ao_nr (c, l, m, n))
    idx_ao = np.argsort (np.asarray (idx_ao))

    if mol.symmetry:
        # I have to figure out what order the Molcas irreps are in on the fly
        # because it seems to change depending on the xyz
        molcas_usymm = molcas_usymm[:,idx_ao]
        proj = np.stack ([(np.dot (molcas_usymm,ir_coeff)**2).sum (1)
            for ir_coeff in mol.symm_orb], axis=0)
        errstr = ("Can't interpret h5 symmetry information; wrong mol? "
                    " <mol.symm_orb|desym_matrix> =\n{}").format (proj)
        assert (np.allclose (np.amax (proj), 1)), errstr
        ix_irrep = np.argmax (proj, axis=0)
        uniq_idx = np.sort (np.unique (ix_irrep, return_index=True)[1])
        uniq_idx = np.append (uniq_idx, len (ix_irrep))
        nmo_irrep = []
        for ix, i in enumerate (uniq_idx[:-1]):
            j = uniq_idx[ix+1]
            assert (len (np.unique (ix_irrep[i:j])) == 1), errstr
            nmo_irrep.append (j-i)
        nmo_irrep = np.asarray (nmo_irrep)
        usymm_irrep_offset = np.cumsum (nmo_irrep) - nmo_irrep
        coeff_irrep_offset = np.cumsum (nmo_irrep**2) - nmo_irrep**2
        mo_coeff = np.zeros ((nao, nao), dtype=np.float_)
        for m_ir, usymm_off, coeff_off in zip (nmo_irrep, usymm_irrep_offset, coeff_irrep_offset):
            i, j = usymm_off, usymm_off+m_ir
            u, v = coeff_off, coeff_off+(m_ir**2)
            usymm = molcas_usymm[i:j,:].T
            coeff = molcas_coeff[u:v].reshape (m_ir, m_ir).T
            mo_coeff[:,i:j] = np.dot (usymm, coeff)
    else:
        assert (molcas_coeff.shape == (nao**2,)), 'mo_vectors.shape = {} but {} AOs'.format (
            molcas_coeff.shape, nao)
        mo_coeff = molcas_coeff.reshape (nao, nao)[:,idx_ao].T
        
    # 'mergesort' keeps degenerate or active-space orbitals in the provided order!
    #  idx_ene = np.argsort (mo_energy, kind='mergesort')

    # modified by Dayou: sort by mo_occ first, then mo_energy
    sort_key = np.min(mo_energy) * 100000 * mo_occ + mo_energy
    idx_ene = np.argsort(sort_key, kind='mergesort')
    mo_coeff = mo_coeff[:,idx_ene]
    mo_occ = mo_occ[idx_ene]
    mo_energy = mo_energy[idx_ene]

    if mol.verbose > logger.INFO:
        fname = str (mol.output)[:-4] + "_h5debug.molden"
        molden.from_mo (mol, fname, mo_coeff, occ=mo_occ, ene=mo_energy)
    if mol.symmetry: # assert (symmetry is right)
        try:
            orbsym = label_orb_symm (mol, mol.irrep_id, mol.symm_orb, mo_coeff)
            mo_coeff = tag_array (mo_coeff, orbsym=orbsym)
        except ValueError as e:
            raise ValueError (("Problem understanding Molcas symmetry?\n"
                "{}").format (str(e)))

    return mo_coeff

def get_mol_from_h5 (h5fname, **kwargs):
    ''' Build a gto.mole object from an h5 file written by OpenMolcas 

        Args:
            h5fname : str
                Path to an h5 file generated by OpenMolcas containing (at least) groups
                'BASIS_FUNCTION_IDS', 'DESYM_MATRIX', 'MO_VECTORS', and 'MO_ENERGIES'
        Kwargs:
            any of the kwargs of gto.M other than atom, basis, symmetry, and unit

        Returns:
            mol: gto.mole object
    '''
    # TODO: somehow preserve OpenMolcas's irrep axes using DESYM_CENTER_LABELS 

    my_symmetry = True
    with h5py.File (h5fname, 'r') as f:
        try:
            symbs = np.asarray (f['DESYM_CENTER_LABELS']).astype ('|U10')
            carts = np.asarray (f['DESYM_CENTER_COORDINATES'])
            basids = np.asarray (f['DESYM_BASIS_FUNCTION_IDS'])
        except KeyError:
            my_symmetry = False
            symbs = np.asarray (f['CENTER_LABELS']).astype ('|U10')
            carts = np.asarray (f['CENTER_COORDINATES'])
            basids = np.asarray (f['BASIS_FUNCTION_IDS'])
        primids = np.asarray (f['PRIMITIVE_IDS'])
        prims = np.asarray (f['PRIMITIVES'])

    # Check for Cartesian bases
    for (c, n, l, m) in basids:
        if l < 0:
            raise NotImplementedError (CARTESIAN)

    # Atomic coordinates and symmetry framework
    symbs = [s.split (' ')[0] for s in symbs]
    my_atom = [[s, xyz.tolist ()] for s, xyz in zip (symbs, carts)]
    if my_symmetry:
        symm, charge_center, axes = detect_symm (my_atom)
        my_symmetry, axes = get_subgroup (symm, axes)
        if my_symmetry.upper () in ('SO3', 'DOOH'):
            my_symmetry = 'D2h'
        if my_symmetry.upper () in ('COOV'):
            my_symmetry = 'C2v'

    # Find basis functions
    my_basis = {}
    for idx, symb in enumerate (symbs):
        if symb in my_basis:
            continue
        thisatom = (primids[:,0] == idx+1)
        atomid = primids[thisatom,1:]
        atomprim = prims[thisatom,:]
        my_basis[symb] = []
        for l in range (np.max (atomid[:,0]+1)):
            thisl = (atomid[:,0] == l)
            lid = atomid[thisl,1:]
            lprim = atomprim[thisl,:]
            thisshl = lid[:,0]==1
            a = lprim[thisshl,0]
            c = lprim[thisshl,1][:,np.newaxis]
            idx_shl = 0
            for n in range (2,np.max (lid)+1):
                thisshl = lid[:,0]==n
                if np.count_nonzero (thisshl) == len (a) and np.all (a == lprim[thisshl,0]):
                    c = np.append (c, lprim[thisshl,1][:,np.newaxis], axis=1)
                else:
                    my_basis[symb].append ([l, np.insert (c, 0, a, axis=1).tolist ()])
                    a = lprim[thisshl,0]
                    c = lprim[thisshl,1][:,np.newaxis]
            my_basis[symb].append ([l, *np.insert (c, 0, a, axis=1).tolist ()])

    for check in ('atom', 'basis', 'symmetry', 'unit'):
        if check in kwargs:
            kwargs.pop (check)
    return gto.M (atom = my_atom, basis = my_basis, symmetry = my_symmetry, unit = 'Bohr', **kwargs)
    #return my_atom, my_basis, my_symmetry


