from __future__ import print_function

"""
Atom Symmetry module

Input: 3D rotation group and a representation of J = [J_{x}, J_{y}, J_{z}]
of an atom

Output:

1) A new representation of Jp and a unitary transformation U such that
   Jp = U^H J U
2) A symbolic matrix representing a general matrix commuting with
   the representation of the group.
"""

import numpy as np
from scipy.sparse import isspmatrix
import sys


def get_3D_rotation_generator():
    return -1.j * np.array([[[int((i - j) * (j - k) * (k - i) / 2)
        for k in range(3)] for j in range(3)] for i in range(3)])


def get_rotation(J, theta, method="eigh"):
    '''
    Get rotation matrix by Lie parameters and generator J.
    '''
    exponent = J[0] * theta[0] + J[1] * theta[1] + J[2] * theta[2]
    if method == "eigh":
        if isspmatrix(exponent):
            exponent = exponent.todense()
        vals, vecs = np.linalg.eigh(exponent)
        rep = np.dot(
            np.dot(vecs, np.diag(np.exp(1.j * vals))), np.conj(vecs).T)
    elif method == "expm":
        from scipy.linalg import expm
        rep = expm(1.j * exponent)
    else:
        raise AssertionError(
                " method =" + method + ": Unsupported in get_rotation!")
    return rep


def get_Lie_parameter_3D(g, plus2pi=False):
    '''
    Get Lie parameter from a 3D rotation matrix.
    '''
    J = get_3D_rotation_generator()
    list_theta = []
    assert len(g) == 3, \
            ' Error: get_Lie_parameter_3D is exclusively for 3D rotations!'

    # all rotation matrices are diagonalizable over the complex field
    vals, vecs = np.linalg.eig(g)

    # pick the label of the eigenvalue 1 (rotation axis)
    rotation_axes_label = np.abs(
        vals - 1.0).tolist().index(min(np.abs(vals - 1.0)))
    axis = vecs[:, rotation_axes_label].real
    fix = (np.trace(g) - 1.0) / 2.0
    if fix <= -1.0:
        assert(np.abs((np.trace(g) - 1.0) / 2.0 - (-1.0)) < 1.e-10)
        fix = -1.0
    if fix >= 1.0:
        assert(np.abs((np.trace(g) - 1.0) / 2.0 - (1.0)) < 1.e-10)
        fix = 1.0
    alpha = np.arccos(fix)
    theta = alpha * axis

    # get the additional sign
    gp = get_rotation(J, theta, method="eigh")
    err = np.max(np.abs(gp - g))
    if err > 1.e-7:
        alpha = -alpha
    list_theta.append(alpha * axis)
    if plus2pi:
        list_theta.append((alpha + np.pi * 2.) * axis)
    return list_theta


def get_Lie_parameters(Rotations, plus2pi=False):
    '''
    Get Lie parameters for a 3D rotation group.
    '''
    Lie_parameters = []
    for g in Rotations:
        list_theta = get_Lie_parameter_3D(g, plus2pi=plus2pi)
        for theta in list_theta:
            Lie_parameters.append(theta)
    return Lie_parameters


def get_plus2pi(Jz):
    '''
    Decide whether one needs to double the number of Lie parameters
    due to half-integer of J.
    '''
    vals = np.linalg.eigvalsh(Jz)
    val = np.min(np.abs(vals))
    if np.abs(val - 0.5) < 1.e-7:
        plus2pi = True
    elif val < 1.e-7:
        plus2pi = False
    else:
        print(' eigen-vales of Jz:\n {}'.format(vals))
        raise AssertionError('Error in eigen-values of Jz!')
    return plus2pi


def check_Lie_parameters(G, J, Lie_parameters):
    '''
    Check the representation generated by Lie parameters and J with G.
    '''
    N = len(Lie_parameters) / len(G)  # can be 1 or 2
    for i, g in enumerate(G):
        for j in range(i * N, (i + 1) * N):
            theta = Lie_parameters[j]
            gp = get_rotation(J, theta, method="eigh")
            err = np.max(np.abs(gp - g))
            assert err < 1.e-7, " check_Lie_parameters!"


def check_Lie_parameters_3D(Rotations, Lie_parameters):
    '''
    Compare the 3D rotation matrices with the one generated
    by Lie parameters and the generator J.
    '''
    J = get_3D_rotation_generator()
    check_Lie_parameters(Rotations, J, Lie_parameters)


def get_ratio_matab(a, b):
    '''
    Get a scale ratio of two matrix.
    '''
    ratio = 0. + 0.j
    # locate the first significant element of a
    ia = -1
    ja = -1
    import itertools
    for i, j in itertools.product(range(len(a)), repeat=2):
        if np.abs(a[i][j]) > 1.e-5:
            ia = i
            ja = j
            break
    if ia >= 0:
        if np.abs(b[ia][ja]) > 1.e-7:
            rat = a[ia, ja] / b[ia, ja]
            if(np.max(np.abs(a - b * rat)) < 1.e-6):
                ratio = rat
    return ratio


def get_structure_constants(J):
    '''
    Get the structure constant or commutation table of [Jx,Jy,Jz].
    '''
    l = len(J)
    eijk = np.zeros((l, l, l), dtype=np.complex)
    for i0, J0 in enumerate(J):
        for i1, J1 in enumerate(J):
            J01 = np.dot(J0, J1) - np.dot(J1, J0)
            for i2, J2 in enumerate(J):
                eijk[i0, i1, i2] = get_ratio_matab(J01, J2)
    return eijk


def check_strcture_constants(J, Jp):
    '''
    Compare the structure constants between two sets of generators.
    '''
    eijk = get_structure_constants(J)
    eijkp = get_structure_constants(Jp)
    if np.max(np.abs(eijkp - eijk)) > 1.e-6:
        print(' Commutation table for J:')
        print(eijk)
        print(' Commutation table for Jp:')
        print(eijkp)
        raise AssertionError('Inconsistency error!')


def check_strcture_constants_3D(J):
    Jp = get_3D_rotation_generator()
    check_strcture_constants(J, Jp)


def get_representation(J, Lie_parameters):
    representation = []
    for theta in Lie_parameters:
        repr = get_rotation(J, theta, method="eigh")
        try:
            # try removing the small elements.
            repr = repr.multiply(np.abs(repr) > 1.e-8)
        except:
            pass
        representation.append(repr)
    return representation


def locate_R_list(R01, R_list):
    '''
    Locate the R in R_list.
    '''
    loc = -1
    min_dist = 1.e5
    for i, R in enumerate(R_list):
        dist = np.max(np.abs(R01 - R))
        if (dist < min_dist):
            min_dist = dist
            if(dist < 1.e-6):
                if loc == -1:
                    loc = i
                else:
                    raise AssertionError(
                            'Found multiple identical elements!')
    if loc == -1:
        print(" Rotation operations list:")
        print_R_list(R_list)
        raise AssertionError(
                "Failed to locate the element with min_dist =" +\
                " %12.3f"%(min_dist))
    return loc


def print_R_list(R_list):
    for i, R in enumerate(R_list):
       print("%3d %5.2f %5.2f %5.2f"%((i,)+tuple(R[0])))
       print("    %5.2f %5.2f %5.2f"%tuple(R[1]))
       print("    %5.2f %5.2f %5.2f"%tuple(R[2]))


def get_product_table(R_list):
    '''
    Get product table of the representation.
    '''
    h = len(R_list)
    table = np.zeros((h, h), dtype=np.int)
    for i0, R0 in enumerate(R_list):
        for i1, R1 in enumerate(R_list):
            R01 = np.dot(R0, R1)
            table[i0, i1] = locate_R_list(R01, R_list)
    return table


def check_product_table(G, representation):
    product_table_G = get_product_table(G)
    product_table_R = get_product_table(representation)
    if len(G) == len(representation):
        err = np.max(np.abs(product_table_G - product_table_R))
        if err > 1.e-6:
            raise AssertionError('Product table not match!')


def get_eigen_space(H, l_list, tol=1.e-6):
    '''
    Get eigen spaces of a Hermitian matrix H,
    in accordance with the possible block form specified by l_list.
    '''
    from scipy.linalg import block_diag
    val_list = []
    vec_list = []
    lbase = 0
    for l1 in l_list:
        vals, vecs = np.linalg.eigh(H[lbase:lbase + l1, lbase:lbase + l1])
        lbase += l1
        val_list += vals.tolist()
        vec_list.append(vecs)
    vals = np.array(val_list)
    vecs = block_diag(*vec_list)
    list_eigen_space = []
    val = vals[0]
    eigen_space = []
    N = len(vals)
    for i in range(N):
        if np.abs((vals[i] - val) / val) < tol:
            eigen_space.append(vecs[:, i])
        else:
            # note eigen_space in c-convention
            list_eigen_space.append(np.array(eigen_space))
            eigen_space = []
            eigen_space.append(vecs[:, i])
            val = vals[i]
    list_eigen_space.append(np.array(eigen_space))
    return list_eigen_space


def brief_eigen_spaces(list_eigen_space):
    print(' Number of eigen space =', len(list_eigen_space))
    print(' Dimension of each eigen space:')
    for i, eigen_space in enumerate(list_eigen_space):
        print(i, len(eigen_space))


def check_commute_G(H, G):
    for i, g in enumerate(G):
        if isinstance(g, np.ndarray):
            C = np.dot(H, g) - np.dot(g, H)
        elif isspmatrix(g):
            C = np.dot(H, g.todense()) - g.dot(H)
        else:
            raise AssertionError("Unsupported datatype of g" + g.__class__)
        err = np.max(np.abs(C))
        if err > 1.e-6:
            raise AssertionError("Error in check_commute_G = %12.7f"%(err))


def get_characters_espace(representation, V):
    '''
    Get the character (vector) of one eigen-space defined by V.
    '''
    chi = []
    rep = []
    for R in representation:
        # note here V in c-convention
        VHRV = V.conj().dot(R).dot(V.T)
        if VHRV.shape == (1,):
            tr = VHRV[0]
        else:
            tr = np.trace(VHRV)
        chi.append(tr)
        rep.append(VHRV)
    return np.array(chi), rep


def get_sub_repr_characters(representation, list_eigen_space):
    list_chi = []
    list_repr = []
    for i, V in enumerate(list_eigen_space):
        chi, rep = get_characters_espace(representation, V)
        list_chi.append(chi)
        list_repr.append(rep)
    return np.array(list_chi), list_repr


def check_sum_chi2_1(chi):
    irr = 'OK'
    h = len(chi)  # order of the group
    sum_chi2 = np.vdot(chi, chi).real
    if np.abs(sum_chi2 - h) > 1.e-5:
        print(' Error in sum_chi2 = ' + str(np.abs(sum_chi2 - h)))
        irr = 'NO'
    return irr


def check_sum_chi2(list_chi):
    '''
    Check sum of character^2
    '''
    for i, chi in enumerate(list_chi):
        irr = check_sum_chi2_1(chi)
        if irr == 'NO':
            return irr
    return irr


def get_ired_representation(representation):
    '''get irreducible representations.
    '''
    l = representation[0].shape[0]
    from pyglib.math.matrix_util import get_matrices_bk_list, sym_array
    from scipy.linalg import block_diag
    check = 'no'
    if '-allmix' in sys.argv:
        l_list = [l]
    else:
        l_list = get_matrices_bk_list(representation)
    print(' list of (sub)dimensions of the representation:\n {}'\
            .format(l_list))

    while check == 'no':
        X = []
        for l1 in l_list:
            X.append(np.random.rand(l1, l1) + np.random.rand(l1, l1) * 1.j)
        X = block_diag(*X)
        # make X Hermitian
        X += X.conj().T
        H = sym_array(X, representation)
        check_commute_G(H, representation)
        list_eigen_space = get_eigen_space(H, l_list)
        list_chi, list_repr = get_sub_repr_characters(
            representation, list_eigen_space)
        check = check_sum_chi2(list_chi)
        if check == 'ok':
            break
    return list_eigen_space, list_chi, list_repr


def get_chi_repr(list_chi, list_repr, list_eigen_space):
    '''
    Group the equivalent (chi) representations.
    '''
    list_chi_repr = []
    list_chi_space = []
    list_distinct_chi = []
    done_list = []
    h = len(list_chi[0])
    chi_product = np.dot(np.conj(list_chi), list_chi.T) / h

    for i, product in enumerate(chi_product):
        chi_repr = []
        chi_space = []
        if i in done_list:
            continue
        for j in range(len(product)):
            if np.abs(product[j] - 1) < 1.e-6:
                if j < i:
                    raise AssertionError('Error in get_chi_repr!')
                chi_repr.append(list_repr[j])
                chi_space.append(list_eigen_space[j])
                done_list.append(j)
                if len(chi_repr) == 1:
                    list_distinct_chi.append(list_chi[j])
        list_chi_repr.append(chi_repr)
        list_chi_space.append(np.asarray(chi_space))
        print( ' chi_space {} dimension = {}'.format(len(list_chi_space) - 1,\
                np.asarray(chi_space).shape))

    return list_chi_repr, list_chi_space, list_distinct_chi


def get_U_equiv_repr(representationA, representationB):
    '''
    Get the transformation matrix U such that B= UAU^H.
    '''
    M = np.zeros(representationA[0].shape)
    for i in range(len(representationA)):
        M = M + representationB[i].dot(representationA[i].T.conj())
    MM = M.dot(M.T.conj())
    vals, vecs = np.linalg.eigh(MM)
    MM = vecs.dot(np.diag(np.sqrt(vals)**(-1))).dot(vecs.T.conj())
    # order is important!
    utrans = MM.dot(M)
    return utrans


def check_U_equiv_repr(representationA, representationB, U):
    for i in range(len(representationA)):
        UAUH = np.dot(U, np.dot(representationA[i], np.conj(U).T))
        err = np.max(np.abs(UAUH - representationB[i]))
        if err > 1.e-5:  # 10:
            raise AssertionError(
                    'Large error in check_U_equiv_repr = %10.5f'%(err))


def get_list_U_equiv_repr(list_chi_repr):
    list_U = []
    for chi_repr in list_chi_repr:
        U = []
        for i in range(1, len(chi_repr)):
            U1 = get_U_equiv_repr(chi_repr[i], chi_repr[0])
            check_U_equiv_repr(chi_repr[i], chi_repr[0], U1)
            U.append(U1)
        list_U.append(U)
    return list_U


def get_U_Jnew(list_chi_space, list_chi_U, reorder=True):
    '''
    Reorder basis to get nice block form.
    '''
    U_Jnew = []
    for i, chi_space in enumerate(list_chi_space):
        for j in range(1, len(chi_space)):
            # Fortran c row/column convention different
            chi_space[j] = np.dot(np.conj(list_chi_U[i][j - 1]), chi_space[j])
        if reorder:
            for k in range(len(chi_space[0])):
                for j in range(len(chi_space)):
                    U_Jnew.append(chi_space[j][k])
        else:
            for j in range(len(chi_space)):
                for k in range(len(chi_space[0])):
                    U_Jnew.append(chi_space[j][k])
    # back to Fortran convention
    U_Jnew = np.array(U_Jnew).T
    print(" U_Jnew unitary error = {}".format(np.max(np.abs(
            U_Jnew.T.conj().dot(U_Jnew)- np.eye(U_Jnew.shape[0])))))
    return U_Jnew


def get_seq_matrix(N):
    '''
    Generate a square matrix with sequential elements.
    '''
    M = []
    for i in range(N):
        Mi = []
        for j in range(N):
            Mi.append(j + 1 + i * N)
        M.append(Mi)
    return np.array(M)


def get_sigma_struct(list_chi_space):
    sigma = []
    ibase = 0
    for chi_space in list_chi_space:
        # number of equivalent irred. repr.
        N = len(chi_space)
        L = len(chi_space[0])
        M = get_seq_matrix(N) + ibase
        for i in range(L):
            sigma.append(M)
        ibase = np.max(M)
    return sigma


def get_Lie_from_JR(Rotations, J):
    '''
    Get the Lie parameters of {R}.
    '''
    plus2pi = get_plus2pi(J[2])
    Lie_parameters = get_Lie_parameters(Rotations, plus2pi=plus2pi)
    check_Lie_parameters_3D(Rotations, Lie_parameters)
    return Lie_parameters


def get_JR_representation(Rotations, J):
    '''
    Get the representation of {R} using generator {J}.
    '''
    Lie_parameters = get_Lie_from_JR(Rotations, J)
    check_strcture_constants_3D(J)
    representation = get_representation(J, Lie_parameters)
    return representation


def get_diag_sub_array(H, I1, I2, J1, J2):
    err = 0.
    N = len(H)
    if I1 > 0:
        err = np.max(err, np.max(np.abs(H[:I1, J1:J2])))
    if I2 < N:
        err = np.max(err, np.max(np.abs(H[I2:, J1:J2])))
    if J1 > 0:
        err = np.max(err, np.max(np.abs(H[I1:I2, :J1])))
    if J2 < N:
        err = np.max(err, np.max(np.abs(H[I1:I2, J2:])))
    assert err < 1.e-7, ' Subarray not diagonal, err > 1.e-7!'
    return H[I1:I2, J1:J2]


def get_atom_U_sigma_chi(representation):
    list_eigen_space, list_chi, list_repr = get_ired_representation(
        representation)
    list_chi_repr, list_chi_space, list_distinct_chi = get_chi_repr(
        list_chi, list_repr, list_eigen_space)
    list_chi_U = get_list_U_equiv_repr(list_chi_repr)
    U = get_U_Jnew(list_chi_space, list_chi_U)
    sigma = get_sigma_struct(list_chi_space)
    return U, sigma, list_distinct_chi


def get_atom_Jnew(Rotations, J):
    '''
    Main function.
    '''
    print(" Rotations list:")
    print_R_list(Rotations)
    representation = get_JR_representation(Rotations, J)
    check_product_table(Rotations, representation)
    U, sigma, chi = get_atom_U_sigma_chi(representation)
    from scipy.linalg import block_diag
    Jnew = []
    for J1 in J:
        Jnew.append(U.conj().T.dot(J1).dot(U))
    return Jnew, U, block_diag(*sigma)


def check_commute_sigma_J(sigma, J, Rotations):
    '''
    Check the commutation of sigma with the new rotated representation.
    '''
    representation = get_JR_representation(Rotations, J)
    check_commute_G(sigma, representation)
