module ProjectionMethods.CollocationMethod

open System
open FsAlg.Generic
open CollocationTaskModel
open Infrastructure

let getChebyshevPolynomialNodes n =
    Array.init n (fun i -> Math.Cos(((2.0 * float i - 1.0) / (2.0 * float n)) * Math.PI))

let getSimpleNodes n =
    let step = 2.0 / (float n - 1.0)
    Array.init n (fun i -> -1.0 + step * float i)

let getRightVector nodes (task : CollocationTask) =
    Vector.init (nodes |> Array.length) (fun i -> nodes[i] |> task.Equation.f)
    
let getCoordinateFunctions n (task : CollocationTask) =
   if n < 2 then raise (invalidArg "n" "n should be more than 2")
   else
       let initialFunctions = [| task.InitialConditions.first_polynomial; task.InitialConditions.second_polynomial |]
       let jacobiPolynomials = (n - 3)
                                  |> getJacobiPolynomials 2
                                  |> Array.map (fun lambda -> (fun x -> x |> lambda |> ((*) ((x * x - 1.0) ** 2.0))))
       Array.concat [ initialFunctions; jacobiPolynomials ]
       
let getMatrix dimension (coordinateFunctions: (float -> float) array) (nodes: float array) (task : CollocationTask) =
    if (coordinateFunctions |> Array.length) <> (nodes |> Array.length) then raise (invalidArg "length" "Coordinate functions count should be equal to the nodes count")
    Matrix.init dimension dimension (fun i j ->
        task.Equation.Ly coordinateFunctions[j] nodes[i])

let getSolutionCoefficients n (coordinateFunctions: (float -> float) array) (nodes: float array) (task: CollocationTask) =
    let matrix = getMatrix n coordinateFunctions nodes task
    let rightVector = getRightVector nodes task
    Matrix.solve matrix rightVector

let getSolution n (coordinateFunctions: (float -> float) array) (task: CollocationTask) =
    let nodes = getSimpleNodes n
    let coefficients = getSolutionCoefficients n coordinateFunctions nodes task

    if (coordinateFunctions |> Array.length) <> (coefficients |> Vector.length) then
        raise (invalidArg "length" "Coordinate functions count should be equal to the coefficients count")

    let rec loop solution i =
        if i = (coordinateFunctions |> Array.length) then
            solution
        else
            loop (fun x -> (x |> solution) + (x |> coordinateFunctions[i]) * coefficients[i]) (i + 1)

    loop (fun _ -> 0.0) 0
    
let getSolutions dimensionArray (task: CollocationTask) =
    let sortedDimensionArray = dimensionArray |> Array.sort

    let coordinateFunctions =
        getCoordinateFunctions (sortedDimensionArray |> Array.last) task

    let rec loop solutions i =
        if i = (sortedDimensionArray |> Array.length) then
            solutions |> List.rev |> List.toArray
        else
            let solution =
                getSolution sortedDimensionArray[i] (Array.sub coordinateFunctions 0 sortedDimensionArray[i]) task

            loop (solution :: solutions) (i + 1)

    loop [] 0
