import scipy
from scipy import optimize
import warnings
warnings.filterwarnings('ignore', 'The iteration is not making good progress')
from Yoast_Glob import *


def funcmain(argList, HEK):
    ORAI1 = HEK[0]
    ORAI2 = HEK[1]
    ORAI3 = HEK[2]
    O1 = argList[0]
    O2 = argList[1]
    O3 = argList[2]
    f1 = 2 * K11 * O1 ** 2 + K12 * O1 * O2 + K13 * O1 * O3 + O1 - I1
    f2 = 2 * K22 * O2 ** 2 + K12 * O1 * O2 + K23 * O2 * O3 + O2 - I2
    f3 = 2 * K33 * O3 ** 2 + K13 * O1 * O3 + K23 * O2 * O3 + O3 - I3
    return [f1, f2, f3]


def func_SKO(argList, SKO):
    SKO_ORAI1 = SKO[0]
    SKO_ORAI2 = SKO[1]
    SKO_ORAI3 = SKO[2]
    O1 = argList [ 0 ]
    O2 = argList [ 1 ]
    O3 = argList [ 2 ]
    if SKO_ORAI1:
        f2 = 2 * K22 * O2 ** 2 + K23 * O2 * O3 + O2 - I2
        f3 = 2 * K33 * O3 ** 2 + K23 * O2 * O3 + O3 - I3
        return [0, f2, f3]
    elif SKO_ORAI2:
        f1 = 2 * K11 * O1 ** 2 + K13 * O1 * O3 + O1 - I1
        f3 = 2 * K33 * O3 ** 2 + K13 * O1 * O3 + O3 - I3
        return [ f1, 0, f3 ]
    elif SKO_ORAI3:
        f1 = 2 * K11 * O1 ** 2 + K12 * O1 * O2 + O1 - I1
        f2 = 2 * K22 * O2 ** 2 + K12 * O1 * O2 + O2 - I2
        return [ f1 ,f2, 0 ]


def func_DKO(argList, DKO ):
    O1 = argList [ 0 ]
    O2 = argList [ 1 ]
    O3 = argList [ 2 ]
    DKO_ORAI1 = DKO[0]
    DKO_ORAI2 = DKO[1]
    DKO_ORAI3 = DKO[2]
    if DKO_ORAI2:
        f2 = 2 * K22 * O2 ** 2 + O2 - I2
        return [0,f2,0]
    elif DKO_ORAI1:
        f1 = 2 * K11 * O1 ** 2 + O1 - I1
        return [f1, 0, 0]
    elif DKO_ORAI3:
        f3 = 2 * K33 * O3 ** 2 + O3 - I3
        return [0, 0, f3]


def func_Jin(result):
    O1 = result [ 0 ]
    O2 = result [ 1 ]
    O3 = result [ 2 ]
    O11 = K11 * (O1 ** 2)
    O22 = K22 * (O2 ** 2)
    O33 = K33 * (O3 ** 2)
    O12 = K12 * O1 * O2
    O23 = K23 * O2 * O3
    O13 = K13 * O1 * O3
    Jin = alpha11 * O11 + alpha12 * O12 + alpha13 * O13 + alpha22 * O22 + alpha23 * O23 + alpha33 * O33
    print('Jin', Jin)
    return Jin


SKO = [0, 1 ,0]
x0 = [0.0, 0, 0.0]
ORAI2_SKO = scipy.optimize.fsolve(func_SKO, x0, args = SKO)
SKO = [0,0,1]
x0 = [0, 0.00005,0]
ORAI3_SKO = scipy.optimize.fsolve(func_SKO, x0, args = SKO)
SKO = [1,0,0]
x0 = [0.0, 0, 0.0]
ORAI1_SKO = scipy.optimize.fsolve(func_SKO, x0, args = SKO)

DKO = [1,0,0]
x0 = [0.0,0,0]
ORAI23_DKO = scipy.optimize.fsolve(func_DKO, x0, args = DKO)

DKO = [0,1,0]
x0 = [0.0,0,0]
ORAI13_DKO = scipy.optimize.fsolve(func_DKO, x0, args = DKO)
DKO = [0,0,0.3]
x0 = [0.,0,0.0]
ORAI12_DKO = scipy.optimize.fsolve(func_DKO, x0, args = DKO)
print('ORAI23_DKO', ORAI23_DKO)

HEK = [1, 1, 1]
x0 = [0, 0, 0.0]
ORAI_HEK = scipy.optimize.fsolve(funcmain, x0, args = HEK)
print('ORAI_HEK', ORAI_HEK)


J_ORAI = {'ORAI_HEK':[], 'ORAI23_DKO':[], 'ORAI13_DKO':[], 'ORAI12_DKO':[], 'ORAI1_SKO':[], 'ORAI2_SKO':[], 'ORAI3_SKO': []}
Jin_ORAI_HEK = func_Jin(ORAI_HEK)
J_ORAI['ORAI_HEK'].append(Jin_ORAI_HEK)
Jin_ORAI23_DKO = func_Jin(ORAI23_DKO)
J_ORAI['ORAI23_DKO'].append(Jin_ORAI23_DKO)

Jin_ORAI13_DKO = func_Jin(ORAI13_DKO)
J_ORAI['ORAI13_DKO'].append(Jin_ORAI13_DKO)
Jin_ORAI12_DKO = func_Jin(ORAI12_DKO)
J_ORAI['ORAI12_DKO'].append(Jin_ORAI12_DKO)
Jin_ORAI1_SKO = func_Jin(ORAI1_SKO)
J_ORAI['ORAI1_SKO'].append(Jin_ORAI1_SKO)
Jin_ORAI2_SKO = func_Jin(ORAI2_SKO)
J_ORAI['ORAI2_SKO'].append(Jin_ORAI2_SKO)
Jin_ORAI3_SKO = func_Jin(ORAI3_SKO)
J_ORAI['ORAI3_SKO'].append(Jin_ORAI3_SKO)
import json
# open file
file = open("ORAI.json", "w")
json.dump(J_ORAI, file)
# close file
file.close()

f = open('ORAI.json', 'r')
if f.mode == 'r':
    ORAI = f.read()
    print(ORAI)