from numba import jit, c16
from numba.experimental import jitclass
from cmath import sqrt
from math import tanh, atanh, pi
from lines import cLineIntersection

def deg2rad(rA):  return rA/180*pi
def rad2degd(rA): return rA*180/pi

@jit(cache=True)
def cap(rN):
    if rN < 0:
        return 0
    return rN
@jit(cache=True)
def ZeroCheck(cN):
    if cN == complex(0,0):
        return complex(1,0)
    else:
        return cN

@jit(cache=True)
def c_tr(cN):
    return cN.real, cN.imag

@jit(cache=True)
def cNorm(cN):
    return ZeroCheck(cN)/abs(ZeroCheck(cN))

@jit(cache=True)
def cDot(cA, cB):
    return (cA * cB.conjugate()).real

@jit(cache=True)
def cDist(cA, cB):
    return abs(cA-cB)

@jit(cache=True)
def MobiusInt(cA,cB,cC,cD): # Bruh
    return Klein2Poincare(cLineIntersection(Poincare2Klein(cA),Poincare2Klein(cB),Poincare2Klein(cC),Poincare2Klein(cD)))

@jit(cache=True)
def MobiusAdd(cA, cB):
    return (cA + cB) / (1 + cA.conjugate() * cB)

@jit(cache=True)
def MobiusGyr(cA, cB):
    return (1 + cA * cB.conjugate()) / (1 + cA.conjugate() * cB)

@jit(cache=True)
def MobiusAddGyr(cA, cB):
    return MobiusAdd(cA, cB), MobiusGyr(cA, cB)

@jit(cache=True)
def MobiusScalar(cA, rT):
    m = abs(cA)
    return tanh(rT * atanh(m)) * cA / m

@jit(cache=True)
def MobiusCoadd(cA,cB):
    return MobiusAdd(cA, (MobiusGyr(cA,-cB) * cB))

@jit(cache=True)
def MobiusCosub(cA,cB):
    return MobiusAdd(cA, -(MobiusGyr(cA,cB) * cB))

@jit(cache=True)
def MobiusLine(cA,cB,rT):
    return MobiusAdd(
            cA,
            MobiusScalar(
                MobiusAdd(-cA,cB),
                rT))

@jit(cache=True)
def MobiusDist(cA,cB):
    return abs(MobiusAdd(-cB,cA))

@jit(cache=True)
def Poincare2Klein(cN):
    return 2*cN / (1 + cDot(cN,cN))

@jit(cache=True)
def Klein2Poincare(cN):
    return (1-sqrt(1-cDot(cN,cN)))*cN/cDot(cN,cN)

gyrosig = [
    ('cPos', c16),
    ('cRot', c16)
]
@jitclass(gyrosig)
class GyroVector:
    def __init__(self, cPos, cRot):
        self.cPos = complex(cPos)
        self.cRot = complex(cRot)
        self.normalize()

    def __add__(gA, gB):
        #cAdd, cGyr = MobiusAddGyr(gA.cPos, gB.cPos / gA.cRot)
        cAdd, cGyr = MobiusAddGyr(gA.cPos, gB.cPos)
        return GyroVector(cAdd, gA.cRot * gB.cRot * cGyr)

    def __neg__(self):
        return GyroVector(-(self.cRot * self.cPos), 1/self.cRot)

    def __sub__(gA, gB):
        return gA + (-gB)

    def copy(self):
        return GyroVector(self.cPos,self.cRot)
    
    def normalize(self):
        if abs(self.cPos) > 1:
            self.cPos = cNorm(self.cPos)
        self.cRot = cNorm(self.cRot)
    
    def rotate(self,cRot):
        self.cRot *= cRot

    def transform(self, cA):
        #cAdd, cGyr = MobiusAddGyr(self.cPos, cA / ZeroCheck(self.cRot))
        cAdd, cGyr = MobiusAddGyr(self.cPos, cA / self.cRot)
        self.cPos = cAdd
        self.cRot = self.cRot * cGyr

    def transformed(self, cA):
        #cAdd, cGyr = MobiusAddGyr(self.cPos, cA / ZeroCheck(self.cRot))
        #return GyroVector(cAdd, self.cRot * cGyr)
        return gA.copy().transform(cA)
    
    def nrtransformed(self, cA):
        return MobiusAdd(self.cPos, cA / self.cRot)

    def dist(self,gA):
        return abs((-self+gA).cPos)