#pragma once

/**
 * @file Solvers.hpp
 *
 * @brief Implémentation de plusieurs algorihtmes de solveurs pour un système
 *        d'équations linéaires
 *
 * Nom: William Nolin
 * Code permanent : NOLW76060101
 * Email : william.nolin.1@ens.etsmtl.ca
 *
 */

#include "Math3D.h"

namespace gti320 {
    // Identification des solveurs
    enum eSolverType {
        kNone, kGaussSeidel, kColorGaussSeidel, kCholesky
    };

    // Paramètres de convergences pour les algorithmes itératifs
    static const float eps = 1e-4f;
    static const float tau = 1e-5f;

    /**
     * Résout Ax = b avec la méthode Gauss-Seidel
     */
    static void gaussSeidel(const Matrix<float, Dynamic, Dynamic> &A,
                            const Vector<float, Dynamic> &b,
                            Vector<float, Dynamic> &x, int k_max) {
        // Implémenter la méthode de Gauss-Seidel
        int n = b.size();

        x.resize(n);
        x.setZero();

        bool converged = false;
        int k = 0;

        do {
            Vector<float, Dynamic> nx = x;

            for (int i = 0; i < n; i++) {
                nx(i) = b(i);

                for (int j = 0; j < i; j++) {
                    nx(i) = nx(i) - A(i, j) * nx(j);
                }

                for (int j = i + 1; j < n; j++) {
                    nx(i) = nx(i) - A(i, j) * x(j);
                }

                nx(i) = nx(i) / A(i, i);
            }

            k++;
            Vector<float, Dynamic> r = A * x - b;

            converged = k >= k_max ||
                        (nx - x).norm() / nx.norm() < tau ||
                        r.norm() / b.norm() < eps;

            x = nx;
        } while (!converged);
    }

    /**

     * Résout Ax = b avec la méthode Gauss-Seidel (coloration de graphe)
     */
    static void gaussSeidelColor(const Matrix<float, Dynamic, Dynamic> &A, const Vector<float, Dynamic> &b,
                                 Vector<float, Dynamic> &x, const Partitions &P, const int maxIter) {
        // Implémenter la méthode de Gauss-Seidel avec coloration de graphe.
        // Les partitions avec l'index de chaque particule sont stockées dans la table des tables, P.
        int n = b.size();

        x.resize(n);
        x.setZero();

        bool converged = false;
        int k = 0;
        Vector<float, Dynamic> nx;

        do {
            nx = x;
            for (const auto &indices: P) {
#pragma omp parallel for
                for (const int &i: indices) {
                    nx(i) = b(i);

                    for (int j = 0; j < i; j++) {
                        nx(i) = nx(i) - A(i, j) * nx(j);
                    }

                    for (int j = i + 1; j < n; j++) {
                        nx(i) = nx(i) - A(i, j) * nx(j);
                    }

                    nx(i) = nx(i) / A(i, i);
                }
            }

            k++;
            Vector<float, Dynamic> r = A * x - b;

            converged = k >= maxIter ||
                        (nx - x).norm() / nx.norm() < tau ||
                        r.norm() / b.norm() < eps;

            x = nx;
        } while (!converged);
    }

    /**
     * Résout Ax = b avec la méthode de Cholesky
     */
    static void cholesky(const Matrix<float, Dynamic, Dynamic> &A,
                         const Vector<float, Dynamic> &b,
                         Vector<float, Dynamic> &x) {
        int n = A.rows();
        x.resize(n);
        x.setZero();

        // Calculer la matrice L de la factorisation de Cholesky
        auto L = Matrix<float, Dynamic, Dynamic>(n, n);

        for (int j = 0; j < n; j++) {
            for (int i = j; i < n; i++) {
                float s = 0;

                for (int k = 0; k < j; k++) {
                    s += L(i, k) * L(j, k);
                }

                if (i == j) {
                    L(i, i) = std::sqrt(A(i, i) - s);
                } else {
                    L(i, j) = (A(i, j) - s) / L(j, j);
                }
            }
        }

        // Résoudre Ly = b
        auto y = Vector<float, Dynamic>(n);
        for (int i = 0; i < n; i++) {
            y(i) = b(i);

            for (int j = 0; j < i; j++) {
                y(i) -= L(i, j) * y(j);
            }

            y(i) /= L(i, i);
        }

        // Résoudre L^t x = y
        // 
        // Remarque : ne pas calculer la transposée de L, c'est inutilement
        // coûteux.
        for (int i = n - 1; i >= 0; i--) {
            x(i) = y(i);

            for (int j = i + 1; j < n; j++) {
                x(i) -= L(j, i) * x(j);
            }

            x(i) /= L(i, i);
        }
    }

}