/* -*- mode: c++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

/*

 Copyright (C) 2002, 2003, 2011 Ferdinando Ametrano

 Copyright (C) 2000, 2001, 2002, 2003 RiskMap srl

 This file is part of QuantLib, a free-software/open-source library

 for financial quantitative analysts and developers - http://quantlib.org/

 QuantLib is free software: you can redistribute it and/or modify it

 under the terms of the QuantLib license.  You should have received a

 copy of the license along with this program; if not, please email

 <quantlib-dev@lists.sf.net>. The license is also available online at

 <https://www.quantlib.org/license.shtml>.

 This program is distributed in the hope that it will be useful, but WITHOUT

 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS

 FOR A PARTICULAR PURPOSE.  See the license for more details.

*/

#include <ql/methods/finitedifferences/tridiagonaloperator.hpp>

namespace QuantLib {

    TridiagonalOperator::TridiagonalOperator(Size size) {

        if (size>=2) {

            n_ = size;

            diagonal_      = Array(size);

            lowerDiagonal_ = Array(size-1);

            upperDiagonal_ = Array(size-1);

            temp_          = Array(size);

        } else if (size==0) {

            n_ = 0;

            diagonal_      = Array(0);

            lowerDiagonal_ = Array(0);

            upperDiagonal_ = Array(0);

            temp_          = Array(0);

        } else {

            QL_FAIL("invalid size (" << size << ") for tridiagonal operator "

                    "(must be null or >= 2)");

        }

    }

    TridiagonalOperator::TridiagonalOperator(const Array& low,

                                             const Array& mid,

                                             const Array& high)

    : n_(mid.size()),

      diagonal_(mid), lowerDiagonal_(low), upperDiagonal_(high), temp_(n_) {

        QL_REQUIRE(low.size() == n_-1,

                   "low diagonal vector of size " << low.size() <<

                   " instead of " << n_-1);

        QL_REQUIRE(high.size() == n_-1,

                   "high diagonal vector of size " << high.size() <<

                   " instead of " << n_-1);

    }

    Array TridiagonalOperator::applyTo(const Array& v) const {

        QL_REQUIRE(n_!=0,

                   "uninitialized TridiagonalOperator");

        QL_REQUIRE(v.size()==n_,

                   "vector of the wrong size " << v.size() <<

                   " instead of " << n_);

        Array result(n_);

        std::transform(diagonal_.begin(), diagonal_.end(),

                       v.begin(),

                       result.begin(),

                       std::multiplies<>());

        // matricial product

        result[0] += upperDiagonal_[0]*v[1];

        for (Size j=1; j<=n_-2; j++)

            result[j] += lowerDiagonal_[j-1]*v[j-1]+

                upperDiagonal_[j]*v[j+1];

        result[n_-1] += lowerDiagonal_[n_-2]*v[n_-2];

        return result;

    }

    Array TridiagonalOperator::solveFor(const Array& rhs) const  {

        Array result(rhs.size());

        solveFor(rhs, result);

        return result;

    }

    void TridiagonalOperator::solveFor(const Array& rhs,

                                       Array& result) const  {

        QL_REQUIRE(n_!=0,

                   "uninitialized TridiagonalOperator");

        QL_REQUIRE(rhs.size()==n_,

                   "rhs vector of size " << rhs.size() <<

                   " instead of " << n_);

        Real bet = diagonal_[0];

        QL_REQUIRE(!close(bet, 0.0),

                   "diagonal's first element (" << bet <<

                   ") cannot be close to zero");

        result[0] = rhs[0]/bet;

        for (Size j=1; j<=n_-1; ++j) {

            temp_[j] = upperDiagonal_[j-1]/bet;

            bet = diagonal_[j]-lowerDiagonal_[j-1]*temp_[j];

            QL_ENSURE(!close(bet, 0.0), "division by zero");

            result[j] = (rhs[j] - lowerDiagonal_[j-1]*result[j-1])/bet;

        }

        // cannot be j>=0 with Size j

        for (Size j=n_-2; j>0; --j)

            result[j] -= temp_[j+1]*result[j+1];

        result[0] -= temp_[1]*result[1];

    }

    Array TridiagonalOperator::SOR(const Array& rhs,

                                   Real tol) const {

        QL_REQUIRE(n_!=0,

                   "uninitialized TridiagonalOperator");

        QL_REQUIRE(rhs.size()==n_,

                   "rhs vector of size " << rhs.size() <<

                   " instead of " << n_);

        // initial guess

        Array result = rhs;

        // solve tridiagonal system with SOR technique

        Real omega = 1.5;

        Real err = 2.0*tol;

        Real temp;

        for (Size sorIteration=0; err>tol ; ++sorIteration) {

            QL_REQUIRE(sorIteration<100000,

                       "tolerance (" << tol << ") not reached in " <<

                       sorIteration << " iterations. " <<

                       "The error still is " << err);

            temp = omega * (rhs[0]     -

                            upperDiagonal_[0]   * result[1]-

                            diagonal_[0]        * result[0])/diagonal_[0];

            err = temp*temp;

            result[0] += temp;

            Size i;

            for (i=1; i<n_-1 ; ++i) {

                temp = omega *(rhs[i]     -

                               upperDiagonal_[i]   * result[i+1]-

                               diagonal_[i]        * result[i] -

                               lowerDiagonal_[i-1] * result[i-1])/diagonal_[i];

                err += temp * temp;

                result[i] += temp;

            }

            temp = omega * (rhs[i]     -

                            diagonal_[i]        * result[i] -

                            lowerDiagonal_[i-1] * result[i-1])/diagonal_[i];

            err += temp*temp;

            result[i] += temp;

        }

        return result;

    }

    TridiagonalOperator TridiagonalOperator::identity(Size size) {

        return TridiagonalOperator(Array(size-1, 0.0),     // lower diagonal

                                   Array(size,   1.0),     // diagonal

                                   Array(size-1, 0.0));    // upper diagonal

    }

}