"""
These functions are mostly deprecated by the C counterparts from lattices.h.
This file is still kept for reference, but might be removed in the future.
"""
import numpy as np
from enum import Enum
from math import floor

nx = None

class LatticeType(Enum):
    """
    All the five Bravais lattices in 2D
    """
    OBLIQUE=1
    RECTANGULAR=2
    SQUARE=4
    RHOMBIC=5
    EQUILATERAL_TRIANGULAR=3
    RIGHT_ISOSCELES=SQUARE
    PARALLELOGRAMMIC=OBLIQUE
    CENTERED_RHOMBIC=RECTANGULAR
    RIGHT_TRIANGULAR=RECTANGULAR
    CENTERED_RECTANGULAR=RHOMBIC
    ISOSCELE_TRIANGULAR=RHOMBIC
    RIGHT_ISOSCELE_TRIANGULAR=SQUARE
    HEXAGONAL=EQUILATERAL_TRIANGULAR

def reduceBasisSingle(b1, b2):
    """
    Lagrange-Gauss reduction of a 2D basis.
    cf. https://www.math.auckland.ac.nz/~sgal018/crypto-book/ch17.pdf
    inputs and outputs are (2,)-shaped numpy arrays
    The output shall satisfy |b1| <= |b2| <= |b2 - b1| 
    TODO doc
    
    TODO perhaps have the (on-demand?) guarantee of obtuse angle between b1, b2?
    TODO possibility of returning the (in-order, no-obtuse angles) b as well?
    """
    b1 = np.array(b1)
    b2 = np.array(b2)
    if b1.shape != (2,) or b2.shape != (2,):
        raise ValueError('Shape of b1 and b2 must be (2,)')
    B1 = np.sum(b1 * b1, axis=-1, keepdims=True)
    mu = np.sum(b1 * b2, axis=-1, keepdims=True) / B1
    b2 = b2 - np.rint(mu) * b1
    B2 = np.sum(b2 * b2, axis=-1, keepdims=True)
    while(np.any(B2 < B1)):
        b2t = b1
        b1 = b2
        b2 = b2t
        B1 = B2
        mu = np.sum(b1 * b2, axis=-1, keepdims=True) / B1
        b2 = b2 - np.rint(mu) * b1
        B2 = np.sum(b2*b2, axis=-1, keepdims=True)
    return np.array((b1,b2))

def shortestBase3(b1, b2):
    ''' 
    returns the "ordered shortest triple" of base vectors (each pair from 
    the triple is a base) and there may not be obtuse angle between b1, b2
    and between b2, b3
    '''
    b1, b2 = reduceBasisSingle(b1,b2)
    if is_obtuse(b1, b2, tolerance=0):
        b3 = b2
        b2 = b2 + b1
    else:
        b3 = b2 - b1
    return (b1, b2, b3)

def shortestBase46(b1, b2, tolerance=1e-13):
    b1, b2 = reduceBasisSingle(b1,b2)
    b1s = np.sum(b1 ** 2)
    b2s = np.sum(b2 ** 2)
    b3 = b2 - b1
    b3s = np.sum(b3 ** 2)
    eps = tolerance * (b2s + b1s)
    if abs(b3s - b2s - b1s) < eps:
        return(b1, b2, -b1, -b2)
    else:
        if b3s - b2s - b1s > eps: #obtuse
            b3 = b2
            b2 = b2 + b1
        return (b1, b2, b3, -b1, -b2, -b3)
    

def is_obtuse(b1, b2, tolerance=1e-13):
    b1s = np.sum(b1 ** 2)
    b2s = np.sum(b2 ** 2)
    b3 = b2 - b1
    b3s = np.sum(b3 ** 2)
    eps = tolerance * (b2s + b1s)
    return (b3s - b2s - b1s > eps)

def classifyLatticeSingle(b1, b2, tolerance=1e-13):
    """
    Given two basis vectors, returns 2D Bravais lattice type.
    Tolerance is relative.
    TODO doc
    """
    b1, b2 = reduceBasisSingle(b1, b2)
    b1s = np.sum(b1 ** 2)
    b2s = np.sum(b2 ** 2)
    b3 = b2 - b1
    b3s = np.sum(b3 ** 2)
    eps = tolerance * (b2s + b1s)
    # Avoid obtuse angle between b1 and b2. TODO This should be yet thoroughly tested.
    # TODO use is_obtuse here?
    if b3s - b2s - b1s > eps:
        b3 = b2
        b2 = b2 + b1
        # N. B. now the assumption |b3| >= |b2| is no longer valid
        #b3 = b2 - b1
        b2s = np.sum(b2 ** 2)
        b3s = np.sum(b3 ** 2)
    if abs(b2s - b1s) < eps or abs(b2s - b3s) < eps: # isoscele
        if abs(b3s - b1s) < eps:
            return LatticeType.EQUILATERAL_TRIANGULAR
        elif abs(b3s - 2 * b1s) < eps:
            return LatticeType.SQUARE
        else:
            return LatticeType.RHOMBIC
    elif abs(b3s - b2s - b1s) < eps:
        return LatticeType.RECTANGULAR
    else:
        return LatticeType.OBLIQUE

def range2D(maxN, mini=1, minj=0, minN = 0):
    """
    "Triangle indices"
    Generates pairs of non-negative integer indices (i, j) such that
    minN ≤ i + j ≤ maxN, i ≥ mini, j ≥ minj.
    TODO doc and possibly different orderings
    """
    for maxn in range(min(mini, minj, minN), floor(maxN+1)): # i + j == maxn
        for i in range(mini, maxn + 1):
            yield (i, maxn - i)

            
def generateLattice(b1, b2, maxlayer=5, include_origin=False, order='leaves'):
    bvs = shortestBase46(b1, b2)
    cc = len(bvs) # "corner count"
    
    if order == 'leaves':
        indices = np.array(list(range2D(maxlayer)))
        ia = indices[:,0]
        ib = indices[:,1]
        cc = len(bvs) # 4 for square/rec,
        leaves = list()
        if include_origin: leaves.append(np.array([[0,0]]))
        for c in range(cc):
            ba = bvs[c]
            bb = bvs[(c+1)%cc]
            leaves.append(ia[:,nx]*ba + ib[:,nx]*bb)
        return np.concatenate(leaves)
    else: 
        raise ValueError('Lattice point order not implemented: ', order)
        
def generateLatticeDisk(b1, b2, r, include_origin=False, order='leaves'):
    b1, b2 = reduceBasisSingle(b1,b2)
    blen = np.linalg.norm(b1, ord=2)
    maxlayer = 2*r/blen # FIXME kanon na vrabce? Nestačí odmocnina ze 2?
    points = generateLattice(b1,b2, maxlayer=maxlayer, include_origin=include_origin, order=order)
    mask = (np.linalg.norm(points, axis=-1, ord=2) <= r)
    return points[mask]

def cellCornersWS(b1, b2,):
    """
    Given basis vectors, returns the corners of the Wigner-Seitz unit cell
    (w1, w2, -w1, w2) for rectangular and square lattice or
    (w1, w2, w3, -w1, -w2, -w3) otherwise
    """
    def solveWS(v1, v2):
        v1x = v1[0]
        v1y = v1[1]
        v2x = v2[0]
        v2y = v2[1]
        lsm = ((-v1y, v2y), (v1x, -v2x))
        rs = ((v1x-v2x)/2, (v1y - v2y)/2)
        t = np.linalg.solve(lsm, rs)
        return np.array(v1)/2 + t[0]*np.array((v1y, -v1x))
    b1, b2 = reduceBasisSingle(b1, b2)
    latticeType = classifyLatticeSingle(b1, b2)
    if latticeType is LatticeType.RECTANGULAR or latticeType is LatticeType.SQUARE:
        return np.array( (
            (+b1+b2),
            (+b2-b1),
            (-b1-b2),
            (-b2+b1),
        )) / 2
    else:
        bvs = shortestBase46(b1,b2,tolerance=0)
        return np.array([solveWS(bvs[i], bvs[(i+1)%6]) for i in range(6)])

def cutWS(points, b1, b2, scale=1., tolerance=1e-13):
    """ 
    From given points, return only those that are inside (or on the edge of)
    the Wigner-Seitz cell of a (scale*b1, scale*b2)-based lattice.
    """
    # TODO check input dimensions?
    bvs = shortestBase46(b1, b2)
    points = np.array(points)
    for b in bvs: 
        mask = (np.tensordot(points, b, axes=(-1, 0)) <=  (scale * (1+tolerance) / 2) *np.linalg.norm(b, ord=2)**2 )
        points = points[mask]
    return points

def filledWS(b1, b2, density=10, scale=1.):
    """
    TODO doc
    TODO more intelligent generation, anisotropy balancing etc.
    """
    points = generateLattice(b1,b2,maxlayer=density*scale, include_origin=True)
    points = cutWS(points/density, np.array(b1)*scale, np.array(b2)*scale)
    return points
    
    
def reciprocalBasis1(*pargs):
    a = reduceBasisSingle(*pargs)
    return np.linalg.inv(a).T

def reciprocalBasis2pi(*pargs):
    return 2*np.pi*reciprocalBasis1(*pargs)

# TODO fill it with "points from reciprocal space" instead
def filledWS2(b1,b2, density=10, scale=1.):
    b1, b2 = reduceBasisSingle(b1,b2)
    b1r, b2r = reciprocalBasis2pi(b1,b2)
    b1l = np.linalg.norm(b1, ord=2)
    b2l = np.linalg.norm(b2, ord=2)
    b1rl = np.linalg.norm(b1r, ord=2)
    b2rl = np.linalg.norm(b2r, ord=2)
    # Black magick. Think later.™ Really. FIXME
    sicher_ratio = np.maximum(b1rl/b2rl, b2rl/b1rl) * np.maximum(b1l/b2l, b2l/b1l) # This really has to be adjusted
    points = generateLattice(b1r,b2r,maxlayer=density*scale*sicher_ratio, include_origin=True)
    points = cutWS(points*b1l/b1rl/density, b1*scale, b2*scale)
    return points


def change_basis(srcbasis, destbasis, srccoords, srccoordsaxis=-1, lattice=True):
    srcbasis = np.array(srcbasis)
    destbasis = np.array(destbasis)
    trmatrix = np.dot(np.linalg.inv(np.transpose(destbasis)), np.transpose(srcbasis))
    if lattice: # if srcbasis and destbasis are two bases of the same lattice, its elements are ints
        otrmatrix = trmatrix
        trmatrix = np.round(trmatrix)
        if not np.all(np.isclose(trmatrix, otrmatrix)):
            raise ValueError("Given srcbasis and destbasis are not bases" 
                "of the same lattice", srcbasis, destbasis, trmatrix-otrmatrix)
    destcoords = np.tensordot(srccoords, trmatrix, axes=(srccoordsaxis, -1))
    return destcoords

"""
TODO
====

- DOC!!!!!
- (nehoří) výhledově pořešit problém „hodně anisotropních“ mřížek (tj. kompensovat
rozdílné délky základních vektorů).

"""