#!/usr/bin/env python
'''
Solver for Model-Order-Reduction with openCFS
'''


import numpy as np
from scipy.io import mmread, mmwrite
import argparse

def reduceSystemMatrix(matrix,modalMatrix):
    reducedMatrix = modalMatrix.T @ matrix @ modalMatrix
    return reducedMatrix

def readComplexMode(modeFilename):
    complexNumbers = np.loadtxt(modeFilename)
    mode = np.array([complexNumbers[:,0] + 1j * complexNumbers[:,1]]).T
    return mode


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Solver for Model-Order-Reduction for harmonic analysis. It recudes a given system in modal coordinates and returns the solution in the coordinates of the given system.")
    parser.add_argument("rhs", help="filename of the right-hand-side vector")
    parser.add_argument("solution", help="filename of the solution-file exported by this program")
    parser.add_argument("modes", help="basename of the file where the eigenmodes are stored")
    parser.add_argument("frequency", help="frequency of the current harmonic step", type=float)
    parser.add_argument("step", help="Current step of the harmonic analysis", type=int)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("-l", "--list", help="determine which modes should be used, e.g. \"1 3 5\" ")
    group.add_argument("-n", "--number", 
                    help="number of modes modes to reduce the system matrix, all modes 1-n will be used",
                    type=int)
    args = parser.parse_args()
    # Angular velocity
    omega = args.frequency*2*np.pi
    # Modes to reduce the system
    if args.list is not None:
        modesList = np.array([int(i) for i in args.list.split()])
    else:
        modesList = np.arange(1,args.number+1)
    
    # If it is the first step the modal matrix must be assembled and saved
    # Afterwards the reduced mass, damping and stiffness matrix as to be loaded or computed
    # The matrices are saved under the following filenames
    modalMatrixFilename = "modalMatrix.mtx"
    reducedMassFilename = "reducedMass.mtx"
    reducedDampFilename = "reducedDamp.mtx"
    reducedStiffFilename = "reducedStiff.mtx"
    reducedRhsFilename = "reducedRhs.mtx"
    if args.step == 1:
        # Assemble the modal matrix
        modalMatrix = readComplexMode(args.modes + "_mode_" + str(modesList[0]) + ".vec")
        for i in range(1,np.size(modesList)):
            mode = readComplexMode(args.modes + "_mode_" + str(modesList[i]) + ".vec")
            modalMatrix = np.append(modalMatrix,mode,axis=1)
        # Save the modal matrix
        mmwrite(modalMatrixFilename,modalMatrix)
        # Load Mass, Damping and Stiffnessmatrix
        massFilename = args.modes + "_mass_0_0.mtx"
        mass = mmread(massFilename)
        try:
            dampingFilename = args.modes + "_damping_0_0.mtx"
            damp = mmread(dampingFilename)
        except:
            print("Warning: No damping matrix found.")
        stiffnessFilename = args.modes + "_stiffness_0_0.mtx"
        stiffness = mmread(stiffnessFilename)
        # Compute and save reduced system matrices
        redMass = reduceSystemMatrix(mass,modalMatrix)
        mmwrite(reducedMassFilename,redMass)
        try:
            redDamp = reduceSystemMatrix(damp,modalMatrix)
            mmwrite(reducedDampFilename,redDamp)
        except:
            redDamp = np.zeros_like(redMass)
        redStiff = reduceSystemMatrix(stiffness,modalMatrix)
        mmwrite(reducedStiffFilename,redStiff)
    else:
        # Load the modal matrix
        modalMatrix = mmread(modalMatrixFilename)
        # Load the reduced system matrices
        redMass = mmread(reducedMassFilename)
        try:
            redDamp = mmread(reducedDampFilename)
        except:
            redDamp = np.zeros_like(redMass)
        redStiff = mmread(reducedStiffFilename)
    # Read the forcing vector and transform it into the modal subspace
    rhs = mmread(args.rhs)
    redRhs = modalMatrix.T @ rhs 
    # Assemble system matrix
    sysMatrix = redStiff+1j*omega*redDamp-omega**2*redMass
    # Compute the solution of the reduced system
    redSolution = np.linalg.solve(sysMatrix,redRhs)
    # Transform the solution back
    solution = modalMatrix @ redSolution
    #solutionVector = modalMatrix @ reducedSolution
    mmwrite(args.solution,solution)