"""
This example program solves the Laplace problem using the Proper Orthogonal
Decomposition (POD) reduced-order modeling technique. For a full description
of the technique the reader is referred to the papers:
[1] K. Kunisch and S. Volkwein. Galerkin proper orthogonal decomposition methods
    for a general equation in fluid dynamics. SIAM Journal on Numerical Analusis,
    40(2):492-515, 2003
[2] S. Volkwein, Optimal control of a phase-field model using the proper orthogonal
    decomposition, Z. Angew. Math. Mech., 81(2001):83-97

The method is split into an offline (computationally intensive) and an online
(computationally cheap) phase. This has many applications including real-time
simulation, uncertainty quantification and inverse problems, where similar
models must be evaluated quickly and many times. 

Offline phase:
    1. A set of solution snapshots of the 1D Laplace problem in the full
       problem space are are constructed and assembled into the columns of a dense
       matrix S.
    2. A standard eigenvalue decomposition is performed on the
       matrix S.T*S.
    3. The eigenvectors and eigenvalues are projected back to the
       original eigenvalue problem S.
    4. The leading eigenvectors then form the POD basis.

Online phase:
    1. The operator corresponding to the discrete Laplacian is
       projected onto the POD basis.
    2. The operator corresponding to the right-hand side is
       projected onto the POD basis.
    3. The reduced (dense) problem expressed in the POD basis is solved.
    4. The reduced solution is projected back to the full
       problem space.

Authors:
Elisa Schenone <elisa.schenone@uni.lu>
Jack S. Hale <jack.hale@uni.lu>
"""

try: range = xrange
except: pass

import sys, slepc4py
slepc4py.init(sys.argv)

from petsc4py import PETSc
from slepc4py import SLEPc
import numpy

import random
import math

def construct_operator(m):
    """
    Standard symmetric eigenproblem corresponding to the Laplacian operator in
    1 dimension with homogeneous Dirichlet boundary conditions.
    """
    # Create matrix for 1D Laplacian operator
    A = PETSc.Mat().create(PETSc.COMM_SELF)
    A.setSizes([m, m])
    A.setFromOptions()
    A.setUp()
    # Fill matrix
    hx = 1.0/(m-1) # x grid spacing
    diagv = 2.0/hx
    offdx = -1.0/hx
    Istart, Iend = A.getOwnershipRange()
    for i in range(Istart, Iend):
        if i != 0 and i != (m - 1):
            A[i, i] = diagv
            if i > 1: A[i, i - 1] = offdx
            if i < m - 2: A[i, i + 1] = offdx
        else:
            A[i, i] = 1.

    A.assemble()

    return A


def set_problem_rhs(m):
    """
    Set bell-shape function as the solution of the Laplacian problem in
    1 dimension with homogeneous Dirichlet boundary conditions and 
    compute the associated discrete RHS.
    """
    # Create 1D mass matrix operator
    M = PETSc.Mat().create(PETSc.COMM_SELF)
    M.setSizes([m, m])
    M.setFromOptions()
    M.setUp()
    # Fill matrix
    hx = 1.0/(m-1) # x grid spacing
    diagv = hx/3
    offdx = hx/6
    Istart, Iend = M.getOwnershipRange()
    for i in range(Istart, Iend):
        if i != 0 and i != (m - 1):
            M[i, i] = 2*diagv
        else:
            M[i, i] = diagv
        if i > 1: M[i, i - 1] = offdx
        if i < m - 2: M[i, i + 1] = offdx
        
    M.assemble()

    x_0 = 0.3
    x_f = 0.7
    mu = x_0 + (x_f - x_0)*random.random()
    sigma = 0.1**2
    uex, f = M.createVecs()
    for j in range(Istart, Iend):
        value = 2/sigma * math.exp(-(hx*j - mu)**2/sigma) * (1 - 2/sigma * (hx*j - mu)**2 )
        f.setValue(j, value)
        value = math.exp(-(hx*j - mu)**2/sigma)
        uex.setValue(j, value)
    f.assemble()
    uex.assemble()

    RHS = f.duplicate()
    M.mult(f, RHS)
    RHS.setValue(0, 0.)
    RHS.setValue(m-1, 0.)
    RHS.assemble()

    return RHS, uex

def solve_laplace_problem(A, RHS):
    """
    Solve 1D Laplace problem with FEM.
    """
    u = A.createVecs('right')
    r, c = A.getOrdering("natural")
    A.factorILU(r, c)
    A.solve(RHS, u)
    A.setUnfactored()
    return u


def solve_laplace_problem_pod(A, RHS, u):
    """ 
    Solve 1D Laplace problem with POD (dense matrix).
    """
    ksp = PETSc.KSP().create(PETSc.COMM_SELF)
    ksp.setOperators(A)
    ksp.setType('preonly')
    pc = ksp.getPC()
    pc.setType('none')
    ksp.setFromOptions()

    ksp.solve(RHS, u)

    return u


def construct_snapshot_matrix(A, N, m):
    """ 
    Set N solution of the 1D Laplace problem as columns of a matrix
    (snapshot matrix). 
    
    Note: For simplicity we do not perform a linear solve, but use
    some analytical solution:
    z(x) = exp(-(x - mu)**2 / sigma)
    """
    snapshots = PETSc.Mat().create(PETSc.COMM_SELF)
    snapshots.setSizes([m, N])
    snapshots.setType('seqdense')
    snapshots.setUp()

    Istart, Iend = snapshots.getOwnershipRange()
    hx = 1.0/(m - 1)
    x_0 = 0.3
    x_f = 0.7
    sigma = 0.1**2
    for i in range(N):
        mu = x_0 + (x_f - x_0)*random.random()
        for j in range(Istart, Iend):
            value = math.exp(-(hx*j - mu)**2/sigma)
            snapshots.setValue(j, i, value)
    snapshots.assemble()

    return snapshots 

def solve_eigenproblem(snapshots, N):
    """
    Solve the eigenvalue problem: the eigenvectors of this problem form the
    POD basis.
    """
    print('Solving POD basis eigenproblem using eigensolver...')
    
    Es = SLEPc.EPS()
    Es.create(PETSc.COMM_SELF)
    Es.setDimensions(N)
    Es.setProblemType(SLEPc.EPS.ProblemType.NHEP) 
    Es.setTolerances(1.0e-8, 500);
    Es.setKrylovSchurRestart(0.6)
    Es.setWhichEigenpairs(SLEPc.EPS.Which.LARGEST_REAL) 
    Es.setOperators(snapshots)
    Es.setFromOptions()

    Es.solve()
    print('Solved POD basis eigenproblem.')
    return Es

def project_STS_eigenvectors_to_S_eigenvectors(bvEs, S):
    sizes = S.getSizes()[0]
    N = bvEs.getActiveColumns()[1]
    bv = SLEPc.BV().create(PETSc.COMM_SELF)
    bv.setSizes(sizes, N)
    bv.setActiveColumns(0, N)
    bv.setFromOptions()
    
    tmpvec2 = S.createVecs('left')
    for i in range(N):
        tmpvec = bvEs.getColumn(i)
        S.mult(tmpvec, tmpvec2)
        bv.insertVec(i, tmpvec2)
        bvEs.restoreColumn(i, tmpvec)

    return bv


def project_reduced_to_full_space(alpha, bv):
    uu = bv.getColumn(0)
    uPOD = uu.duplicate()
    bv.restoreColumn(0,uu)
    
    scatter, Wr = PETSc.Scatter.toAll(alpha)
    scatter.begin(alpha, Wr, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)
    scatter.end(alpha, Wr, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)
    PODcoeff = Wr.getArray(readonly=1)

    bv.multVec(1., 0., uPOD, PODcoeff)

    return uPOD

def main():
    problem_dim = 200
    num_snapshots = 30
    num_pod_basis_functions = 8
    
    assert(num_pod_basis_functions <= num_snapshots)
    
    A = construct_operator(problem_dim)
    S = construct_snapshot_matrix(A, num_snapshots, problem_dim)
    
    # Instead of solving the SVD of S, we solve the standard
    # eigenvalue problem on S.T*S
    STS = S.transposeMatMult(S)
    
    Es = solve_eigenproblem(STS, num_pod_basis_functions)
    nconv = Es.getConverged()
    print('Number of converged eigenvalues: %i' % nconv)
    Es.view()
    
    # get the EPS solution in a BV object
    bvEs = Es.getBV()
    bvEs.setActiveColumns(0, num_pod_basis_functions) 

    # set the bv POD basis
    bv = project_STS_eigenvectors_to_S_eigenvectors(bvEs, S)
    # rescale the eigenvectors
    for i in range(num_pod_basis_functions):
        ll = Es.getEigenvalue(i)
        print('Eigenvalue '+str(i)+': '+str(ll.real))
        bv.scaleColumn(i,1.0/math.sqrt(ll.real))

    print('--------------------------------')
    # Verify that the active columns of bv form an orthonormal subspace, i.e. that X^H*X = Id
    print('Check that bv.dot(bv) is close to the identity matrix')
    XtX = bv.dot(bv)
    XtX.view()
    XtX_array = XtX.getDenseArray()
    n,m = XtX_array.shape
    assert numpy.allclose(XtX_array, numpy.eye(n, m))
    print('--------------------------------')
    print('Solve the problem with POD')

    # Project the linear operator A
    Ared = bv.matProject(A,bv)

    # Set the RHS and the exact solution
    RHS, uex = set_problem_rhs(problem_dim)
    
    # Project the RHS on the POD basis
    RHSred = bv.dotVec(RHS)

    # Solve the problem with POD
    alpha = Ared.createVecs('right')
    alpha = solve_laplace_problem_pod(Ared,RHSred,alpha)

    # Project the POD solution back to the FE space
    uPOD = project_reduced_to_full_space(alpha, bv)

    # Compute the L2 and Linf norm of the error
    error = uex.copy()
    error.axpy(-1,uPOD)
    errorL2 = math.sqrt(error.dot(error).real)
    print('The L2-norm of the error is: '+str(errorL2))

    print("NORMAL END")

if __name__ == '__main__':
    main()