import subprocess
import re
import os
import numpy as np
import matplotlib.pyplot as plt
import math
import fileinput
import time
import matplotlib.ticker as tickers
import matplotlib as mpl
from matplotlib.ticker import MultipleLocator,FormatStrFormatter,MaxNLocator
import codecs
import sys
import threading
# ------------------------------------------------------------------------------------------------------------------------------------------------------------------------
def SetParameterMaterialFunction(inputFunction, parameterName, parameterValue):
    with open(inputFunction+'.py', 'r') as file:
        filedata = file.read()
        filedata = re.sub('(?m)^'+str(parameterName)+'\s?=.*',str(parameterName)+' = '+str(parameterValue),filedata)
        f = open(inputFunction+'.py','w')
        f.write(filedata)
        f.close()


# Read effective quantites
def ReadEffectiveQuantities(QFilePath = os.path.dirname(os.getcwd()) + '/outputs/QMatrix.txt', BFilePath = os.path.dirname(os.getcwd())+ '/outputs/BMatrix.txt'):
    # Read Output Matrices (effective quantities)
    # From Cell-Problem output Files : ../outputs/Qmatrix.txt , ../outputs/Bmatrix.txt
    # -- Read Matrix Qhom
    X = []
    # with codecs.open(path + '/outputs/QMatrix.txt', encoding='utf-8-sig') as f:
    with codecs.open(QFilePath, encoding='utf-8-sig') as f:
        for line in f:
            s = line.split()
            X.append([float(s[i]) for i in range(len(s))])
    Q = np.array([[X[0][2], X[1][2], X[2][2]],
                  [X[3][2], X[4][2], X[5][2]],
                  [X[6][2], X[7][2], X[8][2]] ])

    # -- Read Beff (as Vector)
    X = []
    # with codecs.open(path + '/outputs/BMatrix.txt', encoding='utf-8-sig') as f:
    with codecs.open(BFilePath, encoding='utf-8-sig') as f:
        for line in f:
            s = line.split()
            X.append([float(s[i]) for i in range(len(s))])
    B = np.array([X[0][2], X[1][2], X[2][2]])
    return Q, B

# Function for evaluating the energy in terms of kappa, alpha and Q, B
def eval_energy(kappa,alpha,Q,B)  :
    G=kappa*np.array([[np.cos(alpha)**2],[np.sin(alpha)**2],[np.sqrt(2)*np.cos(alpha)*np.sin(alpha)]])-B
    return np.matmul(np.transpose(G),np.matmul(Q,G))[0,0]
#-------------------------------------------------------------------------------------------------------




# subprocess.call(['python' , 'home/klaus/Desktop/Dune_release/dune-microstructure/experiment/perforated-bilayer/perfBilayer_test.py'])

### DATASET Experiment 1
materialFunctionParameter_1=[
[  # Dataset Ratio r = 0.12
[0.12, 0.0047, 17.32986047, 14.70179844, 0, 1.140351217],
[0.12, 0.0047, 17.32986047, 13.6246,     0, 1.691038688],
[0.12, 0.0047, 17.32986047, 12.42994508, 0, 2.243918105],
[0.12, 0.0047, 17.32986047, 11.69773413, 0, 2.595732726],
[0.12, 0.0047, 17.32986047, 11.14159987, 0, 2.945361006],
[0.12, 0.0047, 17.32986047, 9.500670278, 0, 4.001528043],
[0.12, 0.0047, 17.32986047, 9.005046347, 0, 4.312080261]
],
[  # Dataset Ratio r = 0.17
[0.17, 0.0049, 17.28772791 , 14.75453569, 0, 1.02915975],
[0.17, 0.0049, 17.28772791 , 13.71227639,  0, 1.573720805],
[0.17, 0.0049, 17.28772791 , 12.54975012, 0, 2.407706364],
[0.17, 0.0049, 17.28772791 , 11.83455959, 0, 2.790518802],
[0.17, 0.0049, 17.28772791 , 11.29089521, 0, 3.173814476],
[0.17, 0.0049, 17.28772791 , 9.620608917, 0, 4.187433094],
[0.17, 0.0049, 17.28772791 , 9.101671742, 0, 4.511739072]
],
[  # Dataset Ratio r = 0.22
[0.22, 0.0053,  17.17547062, 14.72680026, 0, 1.058078122],
[0.22, 0.0053,  17.17547062, 13.64338887, 0, 1.544624544],
[0.22, 0.0053,  17.17547062, 12.41305478, 0, 2.317033799],
[0.22, 0.0053,  17.17547062, 11.66482931, 0, 2.686043143],
[0.22, 0.0053,  17.17547062, 11.09781471, 0, 2.967694189],
[0.22, 0.0053,  17.17547062, 9.435795985, 0, 3.913528418],
[0.22, 0.0053,  17.17547062, 8.959564147, 0, 4.262750825]
],
[  # Dataset Ratio r = 0.34
[0.34, 0.0063, 17.14061081 , 14.98380876, 0, 0.789078472],
[0.34, 0.0063, 17.14061081 , 13.97154915,  0, 1.1299263],
[0.34, 0.0063, 17.14061081 , 12.77309253, 0, 1.738136936],
[0.34, 0.0063, 17.14061081 , 12.00959929, 0, 2.159520896],
[0.34, 0.0063, 17.14061081 , 11.42001731, 0, 2.370047499],
[0.34, 0.0063, 17.14061081 , 9.561447179, 0, 3.088299431],
[0.34, 0.0063, 17.14061081 , 8.964704969, 0, 3.18097558]
],
[  # Dataset Ratio r = 0.43
[0.43, 0.0073, 17.07559686 , 15.11316339, 0, 0.577989364],
[0.43, 0.0073, 17.07559686 , 14.17997082, 0, 0.829007544],
[0.43, 0.0073, 17.07559686 , 13.05739844, 0, 1.094211707],
[0.43, 0.0073, 17.07559686 , 12.32309209, 0, 1.325332511],
[0.43, 0.0073, 17.07559686 , 11.74608518, 0, 1.400455154],
[0.43, 0.0073, 17.07559686 , 9.812372466, 0, 1.832325697],
[0.43, 0.0073, 17.07559686 , 9.10519385 , 0, 2.047483977]
],
[  # Dataset Ratio r = 0.49
[0.49, 0.008,  17.01520754, 15.30614414, 0, 0.357615902],
[0.49, 0.008,  17.01520754, 14.49463867, 0, 0.376287785],
[0.49, 0.008,  17.01520754, 13.46629742, 0, 0.851008627],
[0.49, 0.008,  17.01520754, 12.78388234, 0, 0.904475291],
[0.49, 0.008,  17.01520754, 12.23057715, 0, 1.039744708],
[0.49, 0.008,  17.01520754, 10.21852839, 0, 1.346405241],
[0.49, 0.008,  17.01520754, 9.341730605, 0, 1.566568558]
]
]

# ------------------------------------------------------------------------------------------------------------------------------------------------------------------------
### DATASET Experiment 2
materialFunctionParameter_2=[
[  # Dataset Ratio r = 0.12
[0.12, 0.0047, 17.17547062, 8.959564147, 0.0, 0.0 ],
[0.12, 0.0047, 17.17547062, 8.959564147, 0.0, 0.05],
[0.12, 0.0047, 17.17547062, 8.959564147, 0.0, 0.1 ],
[0.12, 0.0047, 17.17547062, 8.959564147, 0.0, 0.15],
[0.12, 0.0047, 17.17547062, 8.959564147, 0.0, 0.2 ],
[0.12, 0.0047, 17.17547062, 8.959564147, 0.0, 0.25],
[0.12, 0.0047, 17.17547062, 8.959564147, 0.0, 0.3 ]
],
[  # Dataset Ratio r = 0.17
[0.17, 0.0049, 17.17547062, 8.959564147, 0.0, 0.0 ],
[0.17, 0.0049, 17.17547062, 8.959564147, 0.0, 0.05],
[0.17, 0.0049, 17.17547062, 8.959564147, 0.0, 0.1 ],
[0.17, 0.0049, 17.17547062, 8.959564147, 0.0, 0.15],
[0.17, 0.0049, 17.17547062, 8.959564147, 0.0, 0.2 ],
[0.17, 0.0049, 17.17547062, 8.959564147, 0.0, 0.25],
[0.17, 0.0049, 17.17547062, 8.959564147, 0.0, 0.3 ]
],
[ # Dataset Ratio r = 0.22
[0.22, 0.0053,  17.17547062, 8.959564147, 0.0, 0.0 ],
[0.22, 0.0053,  17.17547062, 8.959564147, 0.0, 0.05 ],
[0.22, 0.0053,  17.17547062, 8.959564147, 0.0, 0.1 ],
[0.22, 0.0053,  17.17547062, 8.959564147, 0.0, 0.15 ],
[0.22, 0.0053,  17.17547062, 8.959564147, 0.0, 0.2 ],
[0.22, 0.0053,  17.17547062, 8.959564147, 0.0, 0.25 ],
[0.22, 0.0053,  17.17547062, 8.959564147, 0.0, 0.3 ]
],
[  # Dataset Ratio r = 0.34
[0.34, 0.0063, 17.17547062, 8.959564147, 0.0, 0.0 ],
[0.34, 0.0063, 17.17547062, 8.959564147, 0.0, 0.05],
[0.34, 0.0063, 17.17547062, 8.959564147, 0.0, 0.1 ],
[0.34, 0.0063, 17.17547062, 8.959564147, 0.0, 0.15],
[0.34, 0.0063, 17.17547062, 8.959564147, 0.0, 0.2 ],
[0.34, 0.0063, 17.17547062, 8.959564147, 0.0, 0.25],
[0.34, 0.0063, 17.17547062, 8.959564147, 0.0, 0.3 ]
],
[  # Dataset Ratio r = 0.43
[0.43, 0.0073, 17.17547062, 8.959564147, 0.0, 0.0 ],
[0.43, 0.0073, 17.17547062, 8.959564147, 0.0, 0.05],
[0.43, 0.0073, 17.17547062, 8.959564147, 0.0, 0.1 ],
[0.43, 0.0073, 17.17547062, 8.959564147, 0.0, 0.15],
[0.43, 0.0073, 17.17547062, 8.959564147, 0.0, 0.2 ],
[0.43, 0.0073, 17.17547062, 8.959564147, 0.0, 0.25],
[0.43, 0.0073, 17.17547062, 8.959564147, 0.0, 0.3 ]
],
[  # Dataset Ratio r = 0.49
[0.49, 0.008,  17.17547062, 8.959564147, 0.0, 0.0 ],
[0.49, 0.008,  17.17547062, 8.959564147, 0.0, 0.05],
[0.49, 0.008,  17.17547062, 8.959564147, 0.0, 0.1 ],
[0.49, 0.008,  17.17547062, 8.959564147, 0.0, 0.15],
[0.49, 0.008,  17.17547062, 8.959564147, 0.0, 0.2 ],
[0.49, 0.008,  17.17547062, 8.959564147, 0.0, 0.25],
[0.49, 0.008,  17.17547062, 8.959564147, 0.0, 0.3 ]
]
]
# ------------------------------------------------------------------------------------------------------------------------------------------------------------------------
### DATASET Experiment 3
materialFunctionParameter_3=[
[  # Dataset Ratio r = 0.12
[0.12, 0.0047, 17.32986047, 9.005046347, np.pi/12.0, 4.312080261],
[0.12, 0.0047, 17.32986047, 9.005046347, np.pi/6.0, 4.312080261],
[0.12, 0.0047, 17.32986047, 9.005046347, np.pi/4.0, 4.312080261],
[0.12, 0.0047, 17.32986047, 9.005046347, np.pi/3.0, 4.312080261],
[0.12, 0.0047, 17.32986047, 9.005046347, 5.0*np.pi/12.0, 4.312080261],
[0.12, 0.0047, 17.32986047, 9.005046347, np.pi/2.0, 4.312080261],
[0.12, 0.0047, 17.32986047, 9.005046347, 7.0*np.pi/12.0, 4.312080261]
]
]
# ------------------------------------------------------------------------------------------------------------------------------------------------------------------------
# ------------------------------------------------------------------------------------------------------------------------------------------------------------------------
# ------------------------------------------------------------------------------------------------------------------------------------------------------------------------


# 1. material, 2.  material-parameters, 3. ExperimentPathExtension  4. Perforation 5. perforatedLayer  6. Dataset-numbers
scenarios = [["wood_european_beech"   ,   materialFunctionParameter_1  , "/wood-bilayer"              , False , ''       , [0, 1, 2, 3, 4, 5]],    
             ["perforated_wood_upper" ,   materialFunctionParameter_2  , "/perforated-bilayer"        , True  , 'upper'  , [0, 1, 2, 3, 4, 5]],  
             ["perforated_wood_lower" ,   materialFunctionParameter_2  , "/perforated-bilayer"        , True  , 'lower'  , [0, 1, 2, 3, 4, 5]],  
             ["wood_european_beech"   ,   materialFunctionParameter_3  , "/wood-bilayer-rotatedLayer" , False , ''       , [0]]
             ]



print('sys.argv[0]', sys.argv[0])
print('sys.argv[1]', sys.argv[1])
print('sys.argv[2]', sys.argv[2])  
print('sys.argv[3]', sys.argv[3])  
CONTAINER = int(sys.argv[1])


# scenarioNumbers= [int(sys.argv[2]),int(sys.argv[3])]
scenarioNumbers = list(range(int(sys.argv[2]),int(sys.argv[3])+1 ))
print('scenarioNumbers: ', scenarioNumbers)
print('--- Run experiments  ' + str(scenarioNumbers[0]) + ' to ' + str(scenarioNumbers[1]) + ' --- ')

for slurm_array_task_id in scenarioNumbers:

# slurm_array_task_id = int(sys.argv[2])
    print('scenarios[slurm_array_task_id][0]:', scenarios[slurm_array_task_id][0])
    pythonModule = scenarios[slurm_array_task_id][0]
    materialFunctionParameter = scenarios[slurm_array_task_id][1]
    pathExtension = scenarios[slurm_array_task_id][2]
    dataset_numbers  = scenarios[slurm_array_task_id][5]
    perforation = scenarios[slurm_array_task_id][3]
    perforatedLayer = scenarios[slurm_array_task_id][4]

    print('perforation:', perforation)

    #Path for parameterFile
    if CONTAINER:
        #--- Taurus  version
        # print('CONTAINER SETUP USED')
        pythonPath = "/dune/dune-microstructure/experiment/compWood" + pathExtension
        # instrumentedPath = "/dune/dune-gfe/instrumented"
        # resultPath = "/dune/dune-gfe/outputs"
        resultBasePath = "results_"  + scenarios[slurm_array_task_id][0]

        executablePath = "/dune/dune-microstructure/build-cmake/src"
        try:
            os.mkdir(resultBasePath)
        except OSError as error:
            print(error)
    else :
        #--- Local version
        # print('LOCAL SETUP USED')
        pythonPath = "/home/klaus/Desktop/Dune_release/dune-microstructure/experiment/compWood" + pathExtension 
        # instrumentedPath = '/home/klaus/Desktop/harmonicmapBenchmark/dune-gfe/instrumented'
        # resultPath = '/home/klaus/Desktop/harmonicmapBenchmark/dune-gfe/outputs' + "_" + scenarios[slurm_array_task_id][0]
        resultBasePath = '/home/klaus/Desktop/Dune_release/dune-microstructure/experiment/compWood'  + pathExtension 

        executablePath = '/home/klaus/Desktop/Dune_release/dune-microstructure/build-cmake/src'

        try:
            os.mkdir(resultBasePath)

        except OSError as error:
            print(error)


    executable = executablePath + '/Cell-Problem'
    gamma = 1.0


        # path = os.getcwd() + '/experiment/wood-bilayer/results_' + str(dataset_number) + '/'
        # pythonPath = os.getcwd() + '/experiment/wood-bilayer'
        # pythonModule = "wood_european_beech"
        # executable = os.getcwd() + '/build-cmake/src/Cell-Problem'

    # xTest = range(0,np.shape(materialFunctionParameter)[1]);
    # print('range(0,np.shape(materialFunctionParameter)[1]):')

    # for n in xTest:
    #     print(n)

    for dataset_number in dataset_numbers:
        print("------------------")
        print(str(dataset_number) + "th data set")
        print("------------------")
        # ------ Loops through Parameters for Material Function -----------
        for i in range(0,np.shape(materialFunctionParameter)[1]):
            print("------------------")
            print("New Loop")
            print("------------------")
            # Check output directory
            if perforation: 
                outputPath = resultBasePath + '/results_' + perforatedLayer + '_' + str(dataset_number) + '/' + str(i)
                # print('Perforation used')
            else :
                outputPath = resultBasePath + '/results_' + str(dataset_number) + '/' + str(i)
                # print('No Perforation used')
            isExist = os.path.exists(outputPath)
            if not isExist:
                # Create a new directory because it does not exist
                os.makedirs(outputPath)
                print("The new directory " + outputPath + " is created!")

            print('OUTPUTPATH: ', outputPath)

            # thread = threading.Thread(target=run_CellProblem(executable, pythonModule, pythonPath, LOGFILE))
            # thread.start()

            #TODO: apperently its not possible to pass a variable via subprocess and "calculate" another input value inside the python file.
            #      Therefore we use this instead.
            SetParameterMaterialFunction(pythonPath + "/" + pythonModule, "param_r",materialFunctionParameter[dataset_number][i][0])
            SetParameterMaterialFunction(pythonPath + "/" + pythonModule, "param_h",materialFunctionParameter[dataset_number][i][1])
            SetParameterMaterialFunction(pythonPath + "/" + pythonModule, "param_omega_flat",materialFunctionParameter[dataset_number][i][2])
            SetParameterMaterialFunction(pythonPath + "/" + pythonModule, "param_omega_target",materialFunctionParameter[dataset_number][i][3])
            SetParameterMaterialFunction(pythonPath + "/" + pythonModule, "param_theta",materialFunctionParameter[dataset_number][i][4]) 
            if perforation:
                SetParameterMaterialFunction(pythonPath + "/" + pythonModule, "param_beta",materialFunctionParameter[dataset_number][i][5])   


            LOGFILE = outputPath + "/" + pythonModule + "_output" + "_" + str(i) + ".log"

            processList = []
            p = subprocess.Popen(executable + " " + pythonPath + " " + pythonModule
                                            + " -outputPath " + outputPath
                                            + " -gamma " + str(gamma) 
                                            + " | tee " + LOGFILE, shell=True)

            # p = subprocess.Popen(executable + " " + pythonPath + " " + pythonModule
            #                                 + " -outputPath " + outputPath
            #                                 + " -gamma " + str(gamma) 
            #                                 + " -param_r " + str(materialFunctionParameter[i][0])
            #                                 + " -param_h " + str(materialFunctionParameter[i][1])
            #                                 + " -param_omega_flat " + str(materialFunctionParameter[i][2])
            #                                 + " -param_omega_target " + str(materialFunctionParameter[i][3])
            #                                 + " -phase2_angle " + str(materialFunctionParameter[i][4])
            #                                 + " | tee " + LOGFILE, shell=True)

            p.wait() # wait
            processList.append(p)
            exit_codes = [p.wait() for p in processList]
            # ---------------------------------------------------
            # wait here for the result to be available before continuing
            # thread.join()
            f = open(outputPath+"/parameter.txt", "w")
            f.write("r = "+str(materialFunctionParameter[dataset_number][i][0])+"\n")
            f.write("h = "+str(materialFunctionParameter[dataset_number][i][1])+"\n")
            f.write("omega_flat = "+str(materialFunctionParameter[dataset_number][i][2])+"\n")        
            f.write("omega_target = "+str(materialFunctionParameter[dataset_number][i][3])+"\n")         
            f.close()   




print('DONE')