# 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    = ["A0", "A1", "A2", "A3", "A4"]
theta  = [1e2, 1e4, 1e6]
maxrun = 3
part   = ["LD", "MD", "LMD"]
SI     = ['S0', 'S1', 'S2', '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])
rep1    = "n_"+ str(n) + "_k_" + str(nk) + "_" + "EXP3_1"
Wgiven = False
if nk == 1:
    Lliste = [5, 15, 30]
else:
    Lliste = [5, 25,50]
for l2 in Lliste:
    l1 = 0
    for Imat in mat:
        # Build Matrix and EigenPairs :
        wtheta = theta
        if Imat == "A0" or Imat == "A1": wtheta = [1e2] 
        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)
            for ISI in SI:
                if ISI == "S3": wpart = part
                if ISI != "S3": wpart = ["LD"]
                for Ipart in wpart:
                    if Ipart == "LMD":
                        IkListe=[2,4]#nk1 = 2; nk2 = 6; step = 2
                    else:
                        IkListe=[1,2,3,4]#nk1 = 1; nk2 = 5; step = 1
                    if ISI == 'S0' or ISI == 'S1':
                        IkListe=[1]#nk = 1; nk2 = 2; step = 1
                    if ISI == 'S2' : IkListe=[1,2,3,4,10,50,100]#nk1 = 1; nk2 = 7; step = 1
                    irr = "1"; enum = 0; enum_c = 0
                    for Ik in IkListe:#range(nk1, nk2, step):
                        k=Ik; Si=ISI; which=Ipart; RHR="HR"; sigma=1e-7
                        rep = RHR+"_"+rep1
                        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]):
                                if ns > 1  : enum = enum+1; enum_c = enum
                                if ns == 1 : enum_c = 0 
                                Tags = [Imat, str(Irun), Stheta, rep, irr, str(enum_c), ISI, Ipart]
                                s = ConjGrad(A, ritz=True, tol=1e-15, maxiter=n, true_res=False, xexa=xexa[...,i:i+1],
                                             W=W, strat=strat, tags=Tags)
                                irr = "2"
                                if Si != 'S0':
                                    x, W = s.dot(b[...,i:i+1])
                                    csv_eigen(A, W, 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]) )
                                ritz = s.ritz_values()
                                print("Condition number:",ritz[len(ritz)-1:len(ritz)]/ritz[0:1])
    cmd = "R -q --vanilla < EXP3_1.R" + " --args " + rep + " " + RHR
    os.system(cmd)
    cmd = "R -q --vanilla < REXP3_1.R --args "+str(l1)+" "+str(l2) + " " + rep
    os.system(cmd)
cmd = "R -q --vanilla < REXP3_1_L.R" + " --args " + rep
os.system(cmd)