import {Complex} from './Complex.js'
import {maybeComplex, promoteBinary, promoteUnary} from './helpers.js'
import {ResolutionError} from '#core/helpers.js'
import {Returns, ReturnTyping} from '#core/Type.js'
import {match} from '#core/TypePatterns.js'
import {ReturnsAs} from '#generic/helpers.js'

const {conservative, full, free} = ReturnTyping

export const normsq = match(Complex, (math, C, strategy) => {
   const compNormsq = math.normsq.resolve(C.Component, full)
   const R = compNormsq.returns
   const add = math.add.resolve([R,R], strategy)
   return ReturnsAs(add, z => add(compNormsq(z.re), compNormsq(z.im)))
})

export const add = promoteBinary('add')

export const conj = match(Complex, (math, C, strategy) => {
   const neg = math.negate.resolve(C.Component, full)
   const compConj = math.conj.resolve(C.Component, full)
   const cplx = maybeComplex(math, strategy, compConj.returns, neg.returns)
   return ReturnsAs(cplx, z => cplx(compConj(z.re), neg(z.im)))
})

export const divide = [
   match([Complex, T => !T.complex], (math, [C, R], strategy) => {
      const div = math.divide.resolve([C.Component, R], full)
      const cplx = maybeComplex(math, strategy, div.returns, div.returns)
      return ReturnsAs(cplx, (z, r) => cplx(div(z.re, r), div(z.im, r)))
   }),
   match([Complex, Complex], (math, [W, Z], strategy) => {
      const inv = math.invert.resolve(Z, full)
      const mult = math.multiply.resolve([W, inv.returns], strategy)
      return ReturnsAs(mult, (w, z) => mult(w, inv(z)))
   })
]

export const invert = match(Complex, (math, C, strategy) => {
   const conj = math.conj.resolve(C, full)
   const normsq = math.normsq.resolve(C, full)
   const div = math.divide.resolve([C.Component, normsq.returns], full)
   const cplx = maybeComplex(math, strategy, div.returns, div.returns)
   return ReturnsAs(cplx, z => {
      const c = conj(z)
      const d = normsq(z)
      return cplx(div(c.re, d), div(c.im, d))
   })
})

// We want this to work for complex numbers, quaternions, octonions, etc
// See https://math.ucr.edu/home/baez/octonions/node5.html
export const multiply = [
   match([T => !T.complex, Complex], (math, [R, C], strategy) => {
      const mult = math.multiply.resolve([R, C.Component], full)
      const cplx = maybeComplex(math, strategy, mult.returns, mult.returns)
      return ReturnsAs(cplx, (r, z) => cplx(mult(r, z.re), mult(r, z.im)))
   }),
   match([Complex, T => !T.complex], (math, [C, R], strategy) => {
      const mult = math.multiply.resolve([R, C], strategy)
      return ReturnsAs(mult, (z, r) => mult(r, z))
   }),
   match([Complex, Complex], (math, [W, Z], strategy) => {
      const conj = math.conj.resolve(W.Component, full)
      if (conj.returns !== W.Component) {
         throw new ResolutionError(
            `conjugation on ${W.Component} returns type (${conj.returns})`)
      }
      const mWZ = math.multiply.resolve([W.Component, Z.Component], full)
      const mZW = math.multiply.resolve([Z.Component, W.Component], full)
      const sub = math.subtract.resolve([mWZ.returns, mZW.returns], full)
      const add = math.add.resolve([mWZ.returns, mZW.returns], full)
      const cplx = maybeComplex(math, strategy, sub.returns, add.returns)
      return ReturnsAs(cplx, (w, z) => {
         const real = sub(mWZ(     w.re,  z.re), mZW(z.im, conj(w.im)))
         const imag = add(mWZ(conj(w.re), z.im), mZW(z.re,      w.im))
         return cplx(real, imag)
      })
   })
]

export const negate = promoteUnary('negate')

// Should work for complex, quaternions, octonions, etc., even with
// integer coordinates.
export const sqrt = match(Complex, (math, C, strategy) => {
   const re = math.re.resolve(C)
   const R = re.returns
   const isReal = math.isReal.resolve(C)
   // dependencies for the real case:
   const zComp = math.zero(C.Component)
   const sign = math.sign.resolve(R, conservative)
   const oneR = math.one(R)
   const zR = math.zero(R)
   const addRComp = math.add.resolve([R, C.Component], full)
   const sqrtR = math.sqrt.resolve(R, conservative)
   const neg = math.negate.resolve(R, conservative)
   const cplx = math.complex.resolve([C.Component, C.Component], full)
   // additional dependencies for the complex case
   const abs = math.abs.resolve(C, full)
   if (abs.returns !== R) {
      throw new TypeError(`abs on ${C} returns ${abs.returns}, not ${R}`)
   }
   const addRR = math.add.resolve([R, R], conservative)
   const twoR = addRR(oneR, oneR)
   const multRR = math.multiply.resolve([R, R], conservative)
   const im = math.im.resolve(C, full)
   const addRC = math.add.resolve([R, C], full)
   const divRR = math.divide.resolve([R, R], conservative)
   const divCR = math.divide.resolve([C, R], full)

   // The guts of the computation:
   const sqrtImp = Returns(C, z => {
      const a = re(z)
      if (isReal(z)) { // always a special case
         let rp = zComp
         let ip = zComp
         const sgn = sign(a)
         if (sgn === oneR) rp = addRComp(sqrtR(a), rp)
         else if (sgn !== zR) ip = addRComp(sqrtR(neg(a)), ip)
         return cplx(rp, ip)
      }
      // Complex case:
      // We can write z = a + q where q is pure imaginary.
      // Let s = sqrt(2(|z|+a)). Then sqrt(z) = (s/2) + (q/s).
      const s = sqrtR(multRR(twoR, addRR(abs(z), a)))
      const q = im(z)
      return addRC(divRR(s, twoR), divCR(q, s))
   })

   if (strategy != free) return sqrtImp
   const prune = math.pruneImaginary.resolve(C, free)
   return ReturnsAs(prune, z => prune(sqrtImp(z)))
})

export const subtract = promoteBinary('subtract')