module GridMethod.ThomasAlgorithm

open TaskModel

type MatrixRowCoefficients =
    { A: float
      B: float
      C: float
      G: float }

type SolutionComponentCoefficients = { s: float; t: float }

let getMatrixCoefficients (task: Task) grid h =
    Array.init (grid |> Array.length) (fun i ->
        if i = 0 then
            { A = 0.0
              B = h * task.InitialConditions.α_1 + task.InitialConditions.α_2
              C = task.InitialConditions.α_2
              G = -1.0 * h * task.InitialConditions.α }
        else if i = (Array.length grid - 1) then
            { A = task.InitialConditions.β_2
              B = h * task.InitialConditions.β_1 + task.InitialConditions.β_2
              C = 0.0
              G = -1.0 * h * task.InitialConditions.β }
        else
            { A = -1.0 * (grid[i] |> task.Equation.p) - ((grid[i] |> task.Equation.q) * h / 2.0)
              B = -2.0 * (grid[i] |> task.Equation.p) - h * h * (grid[i] |> task.Equation.r)
              C = -1.0 * (grid[i] |> task.Equation.p) + ((grid[i] |> task.Equation.q) * h / 2.0)
              G = h * h * (grid[i] |> task.Equation.f) })

let getSolutionCoefficients (matrixCoefficients: MatrixRowCoefficients array) =
    let rec loop coefficients i =
        if i = (Array.length matrixCoefficients) then
            coefficients
        else if i = 0 then
            let newCoefficient =
                { s = (/) matrixCoefficients[0].C matrixCoefficients[0].B
                  t = (/) (-1.0 * matrixCoefficients[0].G) matrixCoefficients[0].B }

            loop (newCoefficient :: coefficients) (i + 1)
        else
            let newCoefficient =
                { s =
                    (/)
                        matrixCoefficients[i].C
                        (matrixCoefficients[i].B
                         - matrixCoefficients[i].A * (coefficients |> List.head).s)
                  t =
                    (/)
                        (matrixCoefficients[i].A * (coefficients |> List.head).t
                         - matrixCoefficients[i].G)
                        (matrixCoefficients[i].B
                         - matrixCoefficients[i].A * (coefficients |> List.head).s) }

            loop (newCoefficient :: coefficients) (i + 1)

    loop [] 0 |> List.rev |> List.toArray

let getSolution (solutionCoefficients: SolutionComponentCoefficients array) =
    let rec loop solution i =
        if i = (solutionCoefficients |> Array.length) then
            solution
        else if i = 0 then
            loop (solutionCoefficients[Array.length solutionCoefficients - 1].t :: solution) (i + 1)
        else
            let solutionComponentIndex = (-) (Array.length solutionCoefficients - 1) i

            let newComponent =
                (+)
                    ((*) solutionCoefficients[solutionComponentIndex].s (solution |> List.head))
                    solutionCoefficients[solutionComponentIndex].t

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

    loop [] 0 |> List.rev |> List.toArray

let run task grid h =
    getMatrixCoefficients task grid h |> getSolutionCoefficients |> getSolution
