# Testing CG-DEF with W given :

from DDMPY import *
import os
import sys

print("Testing CG-DEF with W=Q.T:")
# Init Test cases :
mat    = ["A2","A3"]
theta  = [1e2, 1e4, 1e6]
maxrun = 1
part   = ["LD","MD"]
SI     = ['S3']
n      = 500
nk     = 1 
if(len(sys.argv) == 1):n = 500
if(len(sys.argv) == 2):n = int(sys.argv[1])
if(len(sys.argv) == 3):
    n  = int(sys.argv[1])
    nk = int(sys.argv[2])
rep    = "n_"+ str(n) + "_k_" + str(nk) + "_" + "EXP1_2"
Wgiven = True
# PERTURBATION:
a = 0.90; b = 1.0;co = (b - a) * np.random.random((400)) + a
a = 1e-6; b = 1e-5;co[0:50] = (b - a) * np.random.random((50)) + a
a = 1e-5; b = 1e-4;co[50:100] = (b - a) * np.random.random((50)) + a
a = 1e-4; b = 1e-3;co[100:150] = (b - a) * np.random.random((50)) + a
a = 1e-3; b = 1e-1;co[150:200] = (b - a) * np.random.random((50)) + a
a = 1e-1; b = 5e-1;co[200:250] = (b - a) * np.random.random((50)) + a
a = 5e-1; b = 0.75;co[250:300] = (b - a) * np.random.random((50)) + a
a = 0.75; b = 0.90;co[300:350] = (b - a) * np.random.random((50)) + a
co = [1e-6,1e-4,1e-2,1e-1,0.25,0.5,0.75,0.85,0.9,0.95,0.999]
pert =  np.random.random((n,1)); npert = la.norm(pert)
for Imat in mat:
    # Build Matrix and EigenPairs :
    wtheta = theta
    if Imat == "A2": a1 = 0 
    if Imat == "A3":
        a1 = n-1
        a = 0.90; b = 1.0;co = (b - a) * np.random.random((16)) + a
        a = 1e-6; b = 1e-5;co[0:2] = (b - a) * np.random.random((2)) + a
        a = 1e-5; b = 1e-4;co[2:4] = (b - a) * np.random.random((2)) + a
        a = 1e-4; b = 1e-3;co[4:6] = (b - a) * np.random.random((2)) + a
        a = 1e-3; b = 1e-1;co[6:8] = (b - a) * np.random.random((2)) + a
        a = 1e-1; b = 5e-1;co[8:10] = (b - a) * np.random.random((2)) + a
        a = 5e-1; b = 0.75;co[10:12] = (b - a) * np.random.random((2)) + a
        a = 0.75; b = 0.90;co[12:14] = (b - a) * np.random.random((2)) + a
        co = 10.0*co
        co = [1e-6,1e-4,1e-2,1e-1,0.25,0.5,0.75,0.85,0.9,0.95,0.999]
    for Itheta in wtheta:
        if Itheta == 1e2: Stheta="1E2"
        if Itheta == 1e4: Stheta="1E4"
        if Itheta == 1e6: Stheta="1E6"
        A,DA,QTA = setmatrix4exp(which=Imat, n=n, k=nk, theta=Itheta)
        irr = "1"
        for ico in co :
            PQTA = QTA[...,a1:a1+1] + ico*(pert/npert)
            NPQTA = la.norm(PQTA); PQTA = PQTA/NPQTA
            CC = np.asscalar(PQTA.T@QTA[...,a1:a1+1])/(la.norm(PQTA)*la.norm(QTA[...,a1:a1+1]))
            SS = np.sqrt(max(0.0,1.0-CC*CC))
            print(ico,npert,la.norm(pert/npert),NPQTA,la.norm(PQTA),CC,SS)
            IkListe = [1]
            for Ik in IkListe:#range(nk1, nk2, step):
                k=Ik; l1=0; l2=10; Si=SI[0]; RHR="Wperturbed"; sigma=1e-7
                if(Imat=="A2"):which=part[0]
                if(Imat=="A3"):which=part[1]
                for Irun in range(1, maxrun+1): 
                # Build exact solutions and right-hand sides :
                    xexa = np.random.rand(A.shape[0],2)
                    b    = A@xexa
                    ns   = 1
                    W    = None
                    strat = [Wgiven, k, l1, l2, Si, which, ns, RHR, sigma, True] 
                    # Solve:
                    for i in range(b.shape[1]):
                        Tags = [Imat, str(Irun), Stheta, rep, irr, str(1), Si, which]
                        # Init W (deflation Subspace) with EigenSpace :
                        if i>0 and Si != 'S0':
                            W = PQTA.copy(); Tags = [Imat, str(Irun), Stheta, rep, irr, str(SS), Si, which]
                        s = ConjGrad(A, ritz=True, tol=1e-16, maxiter=n, true_res=False, xexa=xexa[...,i:i+1],
                                     W=W, strat=strat, tags=Tags)
                        if Si != 'S0':
                            x, W = s.dot(b[...,i:i+1])
                            csv_eigen(A, PQTA.copy(), DA, QTA, strat, tags=Tags)
                        else:
                            x = s.dot(b[...,i:i+1])
                        ns += 1
                        strat = [Wgiven, k, l1, l2, Si, which, ns, RHR, sigma, True]
                        #print("|------> system solved : i=",i)
                        #print(s)
                        #print("||x-xexa||_{A}=",np.sqrt( (x-xexa[...,i:i+1]).T@A@(x-xexa[...,i:i+1]) ) )
                        #print("||b-A@x||_{2}/||b||_{2}",la.norm(b[...,i:i+1]-A@x)/la.norm(b[...,i:i+1]) )
                        #print("Condition number:")
                        #ritz = s.ritz_values()
                        #print(np.amax(ritz)/np.amin(ritz)) 
            irr = "2"
cmd = "R -q --vanilla < EXP1_2.R" + " --args " + rep + " " + RHR
os.system(cmd)