import {Complex, ComplexOp} from './type.js'
import type {
   AbsquareOp, AddOp, AddRealOp, ConjOp, ConservativeSqrtOp, DivideOp,
   DivideByRealOp, MultiplyOp, ReciprocalOp, SqrtOp, SubtractOp,
   UnaryMinusOp
} from '../interfaces/arithmetic.js'
import type {
   NanOp, ReOp, ZeroOp, Depends, RealType, WithConstants, NaNType
} from '../interfaces/type.js'
import type {IsSquareOp, IsRealOp} from '../interfaces/predicate.js'

export const add =
   <T>(dep: Depends<AddOp<T>> & Depends<ComplexOp<T>>): AddOp<Complex<T>> =>
   (w, z) => dep.complex(dep.add(w.re, z.re), dep.add(w.im, z.im))

export const addReal =
   <T>(dep: Depends<AddRealOp<T>> & Depends<ComplexOp<T>>):
      AddRealOp<Complex<T>> =>
   (z, r) => dep.complex(dep.addReal(z.re, r), z.im)

export const unaryMinus =
   <T>(dep: Depends<UnaryMinusOp<T>> & Depends<ComplexOp<T>>):
      UnaryMinusOp<Complex<T>> =>
   z => dep.complex(dep.unaryMinus(z.re), dep.unaryMinus(z.im))

export const conj =
   <T>(dep: Depends<UnaryMinusOp<T>>
       & Depends<ConjOp<T>>
       & Depends<ComplexOp<T>>):
      ConjOp<Complex<T>> =>
   z => dep.complex(dep.conj(z.re), dep.unaryMinus(z.im))

export const subtract =
   <T>(dep: Depends<SubtractOp<T>> & Depends<ComplexOp<T>>):
      SubtractOp<Complex<T>> =>
   (w, z) => dep.complex(dep.subtract(w.re, z.re), dep.subtract(w.im, z.im))

export const multiply =
   <T>(dep: Depends<AddOp<T>>
       & Depends<SubtractOp<T>>
       & Depends<MultiplyOp<T>>
       & Depends<ConjOp<T>>
       & Depends<ComplexOp<T>>):
      MultiplyOp<Complex<T>> =>
   (w, z) => {
      const mult = dep.multiply
      const realpart = dep.subtract(
         mult(         w.re,  z.re), mult(dep.conj(w.im), z.im))
      const imagpart = dep.add(
         mult(dep.conj(w.re), z.im), mult(         w.im,  z.re))
      return dep.complex(realpart, imagpart)
   }

export const absquare =
   <T>(dep: Depends<AddOp<RealType<T>>> & Depends<AbsquareOp<T>>):
      AbsquareOp<Complex<T>> =>
   z => dep.add(dep.absquare(z.re), dep.absquare(z.im))

export const divideByReal =
   <T>(dep: Depends<DivideByRealOp<T>> & Depends<ComplexOp<T>>):
      DivideByRealOp<Complex<T>> =>
   (z, r) => dep.complex(dep.divideByReal(z.re, r), dep.divideByReal(z.im, r))

export const reciprocal =
   <T>(dep: Depends<ConjOp<Complex<T>>>
       & Depends<AbsquareOp<Complex<T>>>
       & Depends<DivideByRealOp<Complex<T>>>):
      ReciprocalOp<Complex<T>> =>
   z => dep.divideByReal(dep.conj(z), dep.absquare(z))

export const divide =
   <T>(dep: Depends<MultiplyOp<Complex<T>>>
       & Depends<ReciprocalOp<Complex<T>>>):
      DivideOp<Complex<T>> =>
   (w, z) => dep.multiply(w, dep.reciprocal(z))

export type ComplexSqrtOp<T> = {
   op?: 'complexSqrt',
   (a: T): Complex<WithConstants<T> | NaNType<WithConstants<T>>>
}
// Complex square root of a real type T
export const complexSqrt =
   <T>(dep: Depends<ConservativeSqrtOp<T>>
       & Depends<IsSquareOp<T>>
       & Depends<UnaryMinusOp<T>>
       & Depends<ComplexOp<WithConstants<T>>>
       & Depends<ZeroOp<T>>
       & Depends<NanOp<Complex<WithConstants<T>>>>): ComplexSqrtOp<T> =>
   r => {
      if (dep.isSquare(r)) return dep.complex(dep.conservativeSqrt(r))
      const negative = dep.unaryMinus(r)
      if (dep.isSquare(negative)) {
         return dep.complex(dep.zero(r), dep.conservativeSqrt(negative))
      }
      // neither the real number or its negative is a square; could happen
      // for example with bigint. So there is no square root. So we have to
      // return the NaN of the type.
      return dep.nan(dep.complex(r))
   }

export const sqrt =
   <T>(dep: Depends<IsRealOp<Complex<T>>>
       & Depends<ComplexSqrtOp<T>>
       & Depends<ConservativeSqrtOp<RealType<T>>>
       & Depends<AbsquareOp<Complex<T>>>
       & Depends<AddRealOp<Complex<T>>>
       & Depends<DivideByRealOp<Complex<T>>>
       & Depends<AddOp<RealType<T>>>
       & Depends<ReOp<Complex<T>>>): SqrtOp<Complex<T>> =>
   z => {
      if (dep.isReal(z)) return dep.complexSqrt(z.re)
      const myabs = dep.conservativeSqrt(dep.absquare(z))
      const num = dep.addReal(z, myabs)
      const r = dep.re(z)
      const denomsq = dep.add(dep.add(myabs, myabs), dep.add(r, r))
      const denom = dep.conservativeSqrt(denomsq)
      return dep.divideByReal(num, denom)
   }

export const conservativeSqrt = sqrt