import {Complex} from './type.js'
import type {
   Dependencies, OpType, OpReturns, RealType, ZeroType
} from '../interfaces/type.js'

declare module "../interfaces/type" {
   interface Operations<T> {
      // TODO: Make Dispatcher collapse operations that match
      // after removing any `_...` suffixes; the following should be
      // additional dispatches for add and divide, not separate
      // operations, in the final mathjs bundle.
      add_real:    {params: [T, RealType<T>], returns: T}
      divide_real: {params: [T, RealType<T>], returns: T}
   }
}

export const add =
   <T>(dep: Dependencies<'add' | 'complex', T>): OpType<'add', Complex<T>> =>
   (w, z) => dep.complex(dep.add(w.re, z.re), dep.add(w.im, z.im))

export const add_real =
   <T>(dep: Dependencies<'add_real' | 'complex', T>):
      OpType<'add_real', Complex<T>> =>
   (z, r) => dep.complex(dep.add_real(z.re, r), z.im)

export const unaryMinus =
   <T>(dep: Dependencies<'unaryMinus' | 'complex', T>):
      OpType<'unaryMinus', Complex<T>> =>
   z => dep.complex(dep.unaryMinus(z.re), dep.unaryMinus(z.im))

export const conj =
   <T>(dep: Dependencies<'unaryMinus' | 'conj' | 'complex', T>):
      OpType<'conj', Complex<T>> =>
   z => dep.complex(dep.conj(z.re), dep.unaryMinus(z.im))

export const subtract =
   <T>(dep: Dependencies<'subtract' | 'complex', T>):
      OpType<'subtract', Complex<T>> =>
   (w, z) => dep.complex(dep.subtract(w.re, z.re), dep.subtract(w.im, z.im))

export const multiply =
   <T>(dep: Dependencies<
       'add' | 'subtract' | 'multiply' | 'conj' | 'complex', T>):
      OpType<'multiply', 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: Dependencies<'absquare', T>
       & Dependencies<'add', OpReturns<'absquare', T>>):
      OpType<'absquare', Complex<T>> =>
   z => dep.add(dep.absquare(z.re), dep.absquare(z.im))

export const divideByReal =
   <T>(dep: Dependencies<'divide_real' | 'complex', T>):
      OpType<'divide_real', Complex<T>> =>
   (z, r) => dep.complex(dep.divide_real(z.re, r), dep.divide_real(z.im, r))

export const reciprocal =
   <T>(dep: Dependencies<'conj' | 'absquare' | 'divide_real', Complex<T>>):
      OpType<'reciprocal', Complex<T>> =>
   z => dep.divide_real(dep.conj(z), dep.absquare(z))

export const divide =
   <T>(dep: Dependencies<'multiply' | 'reciprocal', Complex<T>>):
      OpType<'divide', Complex<T>> =>
   (w, z) => dep.multiply(w, dep.reciprocal(z))

// The dependencies are slightly tricky here, because there are three types
// involved: Complex<T>, T, and RealType<T>, all of which might be different,
// and we have to get it straight which operations we need on each type, and
// in fact, we need `add_real` on both T and Complex<T>, hence the dependency
// with a custom name, not generated via Dependencies<...>
export const sqrt =
   <T>(dep: Dependencies<'add' | 'equal' | 'conservativeSqrt' | 'unaryMinus',
                         RealType<T>>
       & Dependencies<'zero' | 'add_real' | 'complex', T>
       & Dependencies<'absquare' | 're' | 'divide_real', Complex<T>>
       & {add_complex_real: OpType<'add_real', Complex<T>>}):
      OpType<'sqrt', Complex<T>> =>
   z => {
      const myabs = dep.conservativeSqrt(dep.absquare(z))
      const r = dep.re(z)
      const negr = dep.unaryMinus(r)
      if (dep.equal(myabs, negr)) {
         // pure imaginary square root; z.im already zero
         return dep.complex(
            dep.zero(z.re), dep.add_real(z.im, dep.conservativeSqrt(negr)))
      }
      const num = dep.add_complex_real(z, myabs)
      const denomsq = dep.add(dep.add(myabs, myabs), dep.add(r, r))
      const denom = dep.conservativeSqrt(denomsq)
      return dep.divide_real(num, denom)
   }

export const conservativeSqrt = sqrt