import numpy as np
import sys
import os

from .. import handle as hdl

def generators( point_group ):
    """ Provides the symmetry generators for the given point group
    by 3 angles for axis-angle parametrization [omega, theta, phi]

    Parameters
    ----------
    proper_point_group : str
        name of the crystal symmetry, taken from a provided table

    Return values
    ------------
    gen : 2D ndarray, float
        2 symmetry generators in OTP representation
        [[ome1,tta1,phi1], [ome2,tta2,phi2]]
    """
    g1 = [2*np.pi, 0, 0]; # no rotation, any mirror, glide or rotation-inversion

    g2x = [np.pi, np.pi/2, 0.]; # 2nd order rotation around [100]: _2, _2/m
    g2y = [np.pi, np.pi/2, np.pi/2]; # 2nd order rotation around [010]: 2, 2/m
    g2z = [np.pi, 0., 0.]; # 2nd order rotation around [001]: 
    g2d = [np.pi, np.pi/2, np.pi/4]; # 2nd order rotation around [110]

    g3z = [2*np.pi/3, 0., 0.]; # 3rd order rotation around [001]: 3, 3/m
    g3d = [2*np.pi/3, np.arctan(np.sqrt(2)), np.pi/4]; # 3rd order rotation around [111]: _3

    g4z = [np.pi/2, 0., 0.]; # 4th order rotation around [001]: 4, 4/m
    g4x = [np.pi/2, np.pi/2, 0.]; # 4th order rotation around [100]: 4, 4/m

    g6z = [np.pi/3, 0., 0.]; # 3rd order rotation around [001]: 6, 6/m

    match point_group: # last option is the laue group
        case ('1' | 'C1' | '-1'): # 1, m // C1
            gen = np.array((g1,g1),np.float64)
        case ('2' | 'C2' | '2/m'): # 2, 2/m, mm2, -4 // C2
            gen = np.array((g2y,g1),np.float64)
        case ('222' | 'D2' | 'mmm' ): # 222, mmm, -42m // D2
            gen = np.array((g2z,g2x),np.float64)
        case ('4' | 'C4' | '4/m'): # 4, 4/m // C4
            gen = np.array((g4z,g1),np.float64)
        case ('422' | 'D4' | '4/mmm'): # 422, 4/mmm // D4
            gen = np.array((g4z,g2x),np.float64)
        case ('3' | 'C3' | '-3' ): # 3, 3m1, 31m, -3, -6 // C3
            gen = np.array((g3z,g1),np.float64)
        case ('32' | 'D3' | '-3m' ): # 321, 312, -31m, -3m1, -6m2 // D3
            gen = np.array((g3z,g2x),np.float64)
        case ('6' | 'C6' | '6/m'): # 6, 6/m, 6mm // C6
            gen = np.array((g6z,g1),np.float64)
        case ('622' | 'D6' | '6/mmm' ): # 622, 6/mmm // D6
            gen = np.array((g6z,g2x),np.float64)
        case ('23' | 'T' | 'm3' ): # 23, -43m,  m-3 // T
            gen = np.array((g2z,g3d),np.float64)
        case ('432' | 'O' | 'm3m'): # 432, m-3m // O
            gen = np.array((g4z,g4x),np.float64)
        case _:
            print('Point group not recognized')
            sys.exit(1)
    return gen

def cSymmHSH(point_group, n):
    """Loads a matrix to transform HSHs to sHSHs

    loads sets of coefficients that make the hyperspherical 
    harmonic expansion invariant a certain crystal symmetry. 
    These can  be interpreted as defining a compact basis
    for an orientation distribution obeying the required symmetry, 
    namely, the symmetrized hyperspherical harmonics (sHSHs).

    This function loads from files created separately by Matlab
    functions in ressources/symmetrizedHSH

    Parameters
    ----------
    proper_point_group : str
        name of the crystal symmetry, taken from a provided table
    n : int
        order of the HSHs

    Return values
    ------------
    nlm : ndarray, int
        HSH indices corresponding to the coefficients
    c : 2D ndarray, complex
        matrix of HSH coefficients for all sHSHs for given n and lattice
        dimensions: 0: HSH orders, 1: sHSHs
    """
    path_csym = hdl.get_file_path('textom',
            os.path.join('ressources','symmetrizedHSH','output',
                         point_group + '_n' + str(n)))
    # filename = 'ressources/symmetrizedHSH/output/' + point_group + '_n' + str(n)
    if os.path.isfile(path_csym):
        data = np.genfromtxt(
            path_csym,
            dtype=np.complex128,
            skip_header=1,
            skip_footer=0,
        )
        
        nlm = data[:,:3].real.astype(int)
        c = data[:,3:].T
        return nlm, c
    
    else:
        print('Symmetrized HSH base file does not exist. Check for typos or generate it via matlab files')
        sys.exit(1)

def get_NsHSH(point_group, n):
    """ Returns how many sHSHs exist for the respective point group and order

    Parameters
    ----------
    point_group : str
        name of the crystal symmetry, taken from a provided table
    n : int
        order of the HSHs
    """
    if np.mod(n,2):
        return 0
    elif n==0:
        return 1
    else:
        path_overview = hdl.get_file_path('textom',
            os.path.join('ressources','symmetrizedHSH','output','overview.txt'))
        with open(path_overview, "r") as file:
                NsHSH_all = eval( file.read() )[point_group]
        return NsHSH_all[int(n/2-1)]

def get_SFnotation( point_group ):
    """Transforms proper point group from Hermann-Mauguin notation to Schonflies

    This is for connecting to oryx
    """
    convdict = {
        '1': 'C1',
        '2': 'C2',
        '222': 'D2',
        '4': 'C4',
        '422': 'D4',
        '3': 'C3',
        '32': 'D3',
        '6': 'C6',
        '622': 'D6',
        '23': 'T',
        '432': 'O',
        '-1': 'C1',
        '2/m': 'C2',
        'mmmm': 'D2',
        '4/mmm': 'D4',
        '-3m': 'D3',
        '6/mmm': 'D6',
        'm3m': 'O',

    }
    return convdict[point_group]

def get_ppg_notation( diff_laue_group ):
    """Transforms diffraction laue group to corresponding proper point group

    This is for loading the correct symmetrization matrices
    """
    convdict = {
        '-1': '1',
        '2/m': '2',
        'mmmm': '222',
        '4/mmm': '422',
        '-3m': '32',
        '6/mmm': '622',
        'm3m': '432',

    }
    try:
        return convdict[diff_laue_group]
    except:
        return diff_laue_group # this is for the case one provides a ppg

def get_proper_point_group(space_group_number):
    """
    Determine the proper point group from the space group IT number, accounting for equivalent symmetries.

    Parameters:
        space_group_number (int): The IT number of the space group.

    Returns:
        str: The proper point group symbol, or "Unknown" for invalid inputs.
    """
    # consider also:
    # from orix.quaternion.symmetry import get_point_group

    # Define ranges of space groups and their proper point groups (reduced symmetry)
    proper_point_groups = {
        "Triclinic": {range(1, 3): "1"},  # Space groups 1-2
        "Monoclinic": {range(3, 16): "2"},  # Space groups 3-15
        "Orthorhombic": {range(16, 75): "222"},  # Space groups 16-74
        "Tetragonal": {
            range(75, 89): "4",  # Space groups 75-88
            range(89, 143): "422",  # Space groups 89-142
        },
        "Trigonal": {
            range(143, 149): "3",  # Space groups 143-148
            range(149, 168): "32",  # Space groups 149-167
        },
        "Hexagonal": {
            range(168, 177): "6",  # Space groups 168-176
            range(177, 195): "622",  # Space groups 177-194
        },
        "Cubic": {
            range(195, 209): "23",  # Space groups 195-206
            range(209, 231): "432",  # Space groups 207-230
        },
    }

    # Iterate through the dictionary to find the matching range and return the proper point group
    for system, groups in proper_point_groups.items():
        for spg_range, proper_pg in groups.items():
            if space_group_number in spg_range:
                return proper_pg

    return "Unknown"  # For invalid space group numbers

def get_diffraction_laue_group(space_group_number):
    """
    Determine the proper point group from the space group IT number, accounting for equivalent symmetries.

    Parameters:
        space_group_number (int): The IT number of the space group.

    Returns:
        str: The proper point group symbol, or "Unknown" for invalid inputs.
    """
    # consider also:
    # from orix.quaternion.symmetry import get_point_group

    # Define ranges of space groups and their proper point groups (reduced symmetry)
    proper_point_groups = {
        "Triclinic": {range(1, 3): "-1"},  # Space groups 1-2
        "Monoclinic": {range(3, 16): "2/m"},  # Space groups 3-15
        "Orthorhombic": {range(16, 75): "mmm"},  # Space groups 16-74
        "Tetragonal": {
            range(75, 143): "4/mmm",  # Space groups 75-142
        },
        "Trigonal": {
            range(143, 168): "-3m",  # Space groups 143-167
        },
        "Hexagonal": {
            range(168, 195): "6/mmm",  # Space groups 168-194
        },
        "Cubic": {
            range(195, 231): "m3m",  # Space groups 195-206
        },
    }

    # Iterate through the dictionary to find the matching range and return the proper point group
    for system, groups in proper_point_groups.items():
        for spg_range, proper_pg in groups.items():
            if space_group_number in spg_range:
                return proper_pg

    return "Unknown"  # For invalid space group numbers

def zone(proper_point_group, g):
    """Computes the fundamental zone in Rodrigues space for given sample symmetries

    This function constructs a bool-array that can constrict the
    used rotations to the unique ones with respect to symmetries
    in the sample and containing crystals. The corresponding region
    in the rotation space (fundamental zone) is found by converting
    to Rodrigues parametrization and by applying the inequalities
    corresponding to the symmetry, as explained in:
    Heinz, A., and P. Neumann. Acta Crystallographica Section A 
        Foundations of Crystallography 47, no. 6 (November 1, 1991): 
        780–89. https://doi.org/10.1107/S0108767391006864.

    Parameters
    ----------
    proper_point_group : str
        name of the set of rotation symmetries, taken from a provided table
        # add numbering?
    g : 2D ndarray, float
        array of rotations given by angles: [omega, theta, phi]
        where omega is the rotation angle, theta the polar and phi
        the azimutal angle of the rotation axis
        dimensions: 0: rotation index, 1: ome,tta,phi

    Return values
    ------------
    fz : ndarray, bool
        array containing True if the corresponding rotation lies
        in the fundamental zone, False if not
    """
    # transform to Rodrigues vector components
    R1 = np.tan(g[:,0]/2) * np.sin(g[:,1]) * np.cos(g[:,2])
    R2 = np.tan(g[:,0]/2) * np.sin(g[:,1]) * np.sin(g[:,2])
    R3 = np.tan(g[:,0]/2) * np.cos(g[:,1])
    
    match proper_point_group:
        case ('1' | 'C1'): # 1, m // C1
            fz = np.full((g.shape[0]), True)

        case ('2' | 'C2'): # 2, 2/m, mm2, -4 // C2
            fz = np.logical_and.reduce((
                np.abs(R2) <= 1,
            ))
 
        case ('222' | 'D2'): # 222, mmm, -42m // D2
            fz = np.logical_and.reduce((
                np.abs(R1) <= 1,
                np.abs(R2) <= 1,
                np.abs(R3) <= 1,
            ))

        case ('4' | 'C4'): # 4, 4/m // C4
            fz = np.logical_and.reduce((
                np.abs(R1) <= 1,
                np.abs(R2) <= 1,
                +R1+R2 <= np.sqrt(2),
                -R1+R2 <= np.sqrt(2),
                -R1-R2 <= np.sqrt(2),
                +R1-R2 <= np.sqrt(2),
            ))

        case ('422' | 'D4'): # 422, 4/mmm // D4
            fz = np.logical_and.reduce((
                np.abs(R1) <= 1,
                np.abs(R2) <= 1,
                np.abs(R3) <= 1,
                R3 <= np.sqrt(2)-1,
                -R3 <= np.sqrt(2)-1,
                +R1+R2 <= np.sqrt(2),
                -R1+R2 <= np.sqrt(2),
                -R1-R2 <= np.sqrt(2),
                +R1-R2 <= np.sqrt(2),
            ))

        # tan(120deg/2) = sqrt(3)
        case ('3' | 'C3'): # 3, 3m1, 31m, -3, -6 // C3
            a = -1/2 # cos(120)
            b = np.sqrt(3)/2 # sin(120)
            fz = np.logical_and.reduce((       
                +a*R1 +b*R2 <= 1,
                -a*R1 +b*R2 <= 1,
                +a*R1 -b*R2 <= 1,
                -a*R1 -b*R2 <= 1,                
                +b*R1 +a*R2 <= 1,
                -b*R1 +a*R2 <= 1,
                +b*R1 -a*R2 <= 1,
                -b*R1 -a*R2 <= 1,        
            ))

        case ('32' | 'D3'): # 321, 312, -31m, -3m1, -6m2 // D3
            a = -1/2 # cos(120)
            b = np.sqrt(3)/2 # sin(120)
            fz = np.logical_and.reduce((      
                np.abs(R3) <= 1, # 180 deg rotation
                +a*R1 +b*R2 <= 1,
                -a*R1 +b*R2 <= 1,
                +a*R1 -b*R2 <= 1,
                -a*R1 -b*R2 <= 1,                
                +b*R1 +a*R2 <= 1,
                -b*R1 +a*R2 <= 1,
                +b*R1 -a*R2 <= 1,
                -b*R1 -a*R2 <= 1,      
            ))

        case ('6' | 'C6'): # 6, 6/m, 6mm // C6
            a = 1/2 # cos(60)
            b = np.sqrt(3)/2 # sin(60)
            fz = np.logical_and.reduce((
                np.abs(R1) <= 1, # 180 deg rotation
                np.abs(R2) <= 1, # 180 deg rotation
                +a*R1 +b*R2 <= 1,
                -a*R1 +b*R2 <= 1,
                +a*R1 -b*R2 <= 1,
                -a*R1 -b*R2 <= 1,                
                +b*R1 +a*R2 <= 1,
                -b*R1 +a*R2 <= 1,
                +b*R1 -a*R2 <= 1,
                -b*R1 -a*R2 <= 1,          
            ))

        case ('622' | 'D6'): # 622, 6/mmm // D6
            a = 1/2
            b = np.sqrt(3)/2
            fz = np.logical_and.reduce((
                np.abs(R1) <= 1, # 180 deg rotation
                np.abs(R2) <= 1, # 180 deg rotation
                np.abs(R3) <= 1, # 180 deg rotation
                +a*R1 +b*R2 <= 1,
                -a*R1 +b*R2 <= 1,
                +a*R1 -b*R2 <= 1,
                -a*R1 -b*R2 <= 1,                
                +b*R1 +a*R2 <= 1,
                -b*R1 +a*R2 <= 1,
                +b*R1 -a*R2 <= 1,
                -b*R1 -a*R2 <= 1,
                +a +b*R3 <= 1,
                -a +b*R3 <= 1,
                +a -b*R3 <= 1,
                -a -b*R3 <= 1,
                +b +a*R3 <= 1,
                -b +a*R3 <= 1,
                +b -a*R3 <= 1,
                -b -a*R3 <= 1,               
            ))

        case ('23' | 'T'): # 23, -43m,  m-3 // T
            fz = np.logical_and.reduce((
                # missing ############## ?
                R1 <= np.sqrt(2)-1,
                -R1 <= np.sqrt(2)-1,
                R2 <= np.sqrt(2)-1,
                -R2 <= np.sqrt(2)-1,
                R3 <= np.sqrt(2)-1,
                -R3 <= np.sqrt(2)-1,
            ))

        case ('432' | 'O'): # 432, m-3m // O
            fz = np.logical_and.reduce((
                R1 <= np.sqrt(2)-1,
                -R1 <= np.sqrt(2)-1,
                R2 <= np.sqrt(2)-1,
                -R2 <= np.sqrt(2)-1,
                R3 <= np.sqrt(2)-1,
                -R3 <= np.sqrt(2)-1,
                +R1+R2+R3 <= 1,
                -R1+R2+R3 <= 1,
                -R1-R2+R3 <= 1,
                -R1-R2-R3 <= 1,
                +R1-R2+R3 <= 1,
                +R1-R2-R3 <= 1,
                +R1+R2-R3 <= 1,
            ))
            
        case _:
            print('Point group not recognized')
            sys.exit(1)
    return fz

def miller_to_hex(indices):
    """
    Convert 3-index Miller indices (h,k,l) to 4-index hexagonal (h,k,i,l).
    
    Parameters
    ----------
    indices : array-like, shape (3,) or (N,3)
        Input Miller indices (h,k,l).
    
    Returns
    -------
    hex_indices : ndarray, shape (4,) or (N,4)
        Converted hexagonal indices (h,k,i,l).
    """
    arr = np.atleast_2d(indices).astype(int)
    h, k, l = arr[:, 0], arr[:, 1], arr[:, 2]
    i = -(h + k)
    hex_arr = np.stack([h, k, i, l], axis=1)
    return hex_arr if indices.ndim > 1 else hex_arr[0]


def hex_to_miller(indices, check=True):
    """
    Convert 4-index hexagonal (h,k,i,l) to 3-index Miller (h,k,l).
    
    Parameters
    ----------
    indices : array-like, shape (4,) or (N,4)
        Input hexagonal indices (h,k,i,l).
    check : bool
        If True, enforce i = -(h+k).
    
    Returns
    -------
    miller_indices : ndarray, shape (3,) or (N,3)
        Converted Miller indices (h,k,l).
    """
    arr = np.atleast_2d(indices).astype(int)
    h, k, i, l = arr[:, 0], arr[:, 1], arr[:, 2], arr[:, 3]

    if check:
        expected = -(h + k)
        if not np.all(i == expected):
            raise ValueError(f"Inconsistent indices: expected i={expected}, got {i}")

    miller_arr = np.stack([h, k, l], axis=1)
    return miller_arr if indices.ndim > 1 else miller_arr[0]