module GridMethod.Algorithm

open TaskModel
open Infrastructure
open FsAlg.Generic

let getInitialConditionsFilledMatrix x_points t_points (task: Task) =
    Matrix.init (t_points |> Array.length) (x_points |> Array.length) (fun i j ->
        if i = 0 then
            task.InitialConditions.µ x_points[j]
        else if j = 0 then
            task.InitialConditions.µ_1 t_points[i]
        else if j = (x_points |> Array.length) then
            task.InitialConditions.µ_2 t_points[i]
        else
            0.0)

let getAuxiliaryMatrix dimension h t (task: Task) =
    Matrix.init dimension dimension (fun i j ->
        if i = 0 && j = 0 then
            1.0
        else if i = dimension - 1 && j = dimension - 1 then
            1.0
        else if i > 0 && i < dimension - 1 then
            if j = i - 1 || j = i + 1 then
                -1.0 * t * task.Equation.κ / (h ** 2.0)
            else if j = i then
                2.0 * t * task.Equation.κ / (h ** 2.0) + 1.0
            else
                0.0
        else
            0.0)

let runWithExplicitScheme (task: Task) x_segments_count time_segments_count =
    let x_points = getIntervalPartition 0.0 task.InitialConditions.a x_segments_count
    let h = getStep 0.0 task.InitialConditions.a x_segments_count

    let t_points = getIntervalPartition 0.0 task.InitialConditions.T time_segments_count
    let t = getStep 0.0 task.InitialConditions.T time_segments_count

    let mutable solution = getInitialConditionsFilledMatrix x_points t_points task

    for i in 1..time_segments_count do
        for j in 1 .. x_segments_count - 1 do
            solution[i, j] <-
                solution[i - 1, j]
                + t * task.Equation.f t_points[i - 1] x_points[j]
                + (t * task.Equation.κ / (h ** 2.0))
                  * (solution[i - 1, j - 1] - 2.0 * solution[i - 1, j] + solution[i - 1, j + 1])

    { Solution = solution
      t_points = t_points
      x_points = x_points }

let runWithImplicitScheme (task: Task) x_segments_count time_segments_count =
    let x_points = getIntervalPartition 0.0 task.InitialConditions.a x_segments_count
    let h = getStep 0.0 task.InitialConditions.a x_segments_count

    let t_points = getIntervalPartition 0.0 task.InitialConditions.T time_segments_count
    let t = getStep 0.0 task.InitialConditions.T time_segments_count

    let mutable solution = getInitialConditionsFilledMatrix x_points t_points task

    let auxiliaryMatrix = getAuxiliaryMatrix (x_points |> Array.length) h t task

    for i in 1..time_segments_count do
        let auxiliaryRightVector =
            Vector.init (x_points |> Array.length) (fun k ->
                if k = 0 then
                    task.InitialConditions.µ_1 t_points[i]
                else if k = x_segments_count then
                    task.InitialConditions.µ_2 t_points[i]
                else
                    t * task.Equation.f t_points[i] x_points[k] + solution[i - 1, k])

        let localSolution = Matrix.solve auxiliaryMatrix auxiliaryRightVector

        for j = 0 to (Array.length x_points - 1) do
            solution[i, j] <- localSolution[j]

    { Solution = solution
      t_points = t_points
      x_points = x_points }
