module ProjectionMethods.RitzMethod

open Infrastructure
open MathNet.Numerics
open MathNet.Numerics.Differentiate
open RitzTaskModel
open FsAlg.Generic

let calculate_Q_l (task: RitzTask) y z =
    if
        abs (task.InitialConditions.α_1 - 0.0) < 1e-7
        || abs (task.InitialConditions.α_2 - 0.0) < 1e-7
    then
        0.0
    else
        ((/) task.InitialConditions.α_1 task.InitialConditions.α_2)
        * (-1.0 |> task.Equation.p)
        * (-1.0 |> y)
        * (-1.0 |> z)

let calculate_Q_r (task: RitzTask) y z =
    if
        abs (task.InitialConditions.β_1 - 0.0) < 1e-7
        || abs (task.InitialConditions.β_2 - 0.0) < 1e-7
    then
        0.0
    else
        ((/) task.InitialConditions.β_1 task.InitialConditions.β_2)
        * (1.0 |> task.Equation.p)
        * (1.0 |> y)
        * (1.0 |> z)

let calculateBilinearForm (task: RitzTask) y z =
    let integrand =
        fun x ->
            (+)
                ((x |> task.Equation.p) * (x |> derivativeFunc 1 y) * (x |> derivativeFunc 1 z))
                ((x |> task.Equation.r) * (x |> y) * (x |> z))

    let Q_l = calculate_Q_l task y z
    let Q_r = calculate_Q_r task y z
    Integrate.OnClosedInterval(integrand, -1.0, 1.0) |> (+) Q_l |> (+) Q_r

let getCoordinateFunctions n =
    (n - 1)
    |> getJacobiPolynomials 1
    |> Array.map (fun lambda -> (fun x -> x |> lambda |> ((*) (1.0 - x * x))))

let getCoordinateFunctionDerivatives n =
    n |> getCoordinateFunctions |> Array.map (derivativeFunc 1)

let getMatrix dimension (coordinateFunctions: (float -> float) array) (task: RitzTask) =
    Matrix.init dimension dimension (fun i j ->
        (calculateBilinearForm task coordinateFunctions[j] coordinateFunctions[i]))

let getRightVector dimension (coordinateFunctions: (float -> float) array) (task: RitzTask) =
    Vector.init dimension (fun i ->
        let integrand = fun x -> (x |> task.Equation.f) * (x |> coordinateFunctions[i])
        Integrate.OnClosedInterval(integrand, -1.0, 1.0))

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

let getSolution n (coordinateFunctions: (float -> float) array) (task: RitzTask) =
    let coefficients = getSolutionCoefficients n coordinateFunctions 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: RitzTask) =
    let sortedDimensionArray = dimensionArray |> Array.sort

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

    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
