import numpy as np
import matplotlib.pyplot as plt
import mapp4py
from mapp4py import md
from lib.elasticity import rot, cubic, resize, displace, crack

"""
Explanation: Introduction
Preperation
Import packages
End of explanation
"""


from mapp4py import mpi
if mpi().rank!=0:
    with open(os.devnull, 'w') as f:
        sys.stdout = f;

"""
Explanation: Block the output of all cores except for one
End of explanation
"""


xprt = md.export_cfg("");

"""
Explanation: Define an md.export_cfg object
md.export_cfg has a call method that we can use to create quick snapshots of our simulation box
End of explanation
"""


_ = np.array([[-1,1,0],[1,1,1],[1,1,-2]], dtype=np.float);
Q = np.linalg.inv(np.sqrt(_ @ _.T)) @ _;

C = rot(cubic(1.3967587463636366,0.787341583191591,0.609615090769241),Q)

B = np.linalg.inv(
    np.array([
    [C[0, 0, 0, 0], C[0, 0, 1, 1], C[0, 0, 0, 1]],
    [C[0, 0, 1, 1], C[1, 1, 1, 1], C[1, 1, 0, 1]],
    [C[0, 0, 0, 1], C[1, 1, 0, 1], C[0, 1, 0, 1]]
    ]
))

_ = np.roots([B[0, 0], -2.0*B[0, 2],2.0*B[0, 1]+B[2, 2], -2.0*B[1, 2], B[1, 1]])

mu = np.array([_[0],0.0]);

if np.absolute(np.conjugate(mu[0]) - _[1]) > 1.0e-12:
    mu[1] = _[1];
else:
    mu[1] = _[2]

alpha = np.real(mu);
beta = np.imag(mu);

p = B[0,0] * mu**2 - B[0,2] * mu + B[0, 1]
q = B[0,1] * mu - B[0, 2] + B[1, 1]/ mu

K = np.stack([p, q]) * np.array(mu[1], mu[0]) /(mu[1] - mu[0])

K_r = np.real(K)
K_i = np.imag(K)

Tr = np.stack([
    np.array(np.array([[1.0, alpha[0]], [0.0, beta[0]]])), 
    np.array([[1.0, alpha[1]], [0.0, beta[1]]])
], axis=1)


def u_f0(x): return np.sqrt(np.sqrt(x[0] * x[0] + x[1] * x[1]) + x[0])
def u_f1(x): return np.sqrt(np.sqrt(x[0] * x[0] + x[1] * x[1]) - x[0]) * np.sign(x[1]) 


def disp(x): 
    _ = Tr @ x
    return K_r @ u_f0(_) + K_i @ u_f1(_)

n = 300;
r = 10;
disp_scale = 0.3;

n0 = int(np.round(n/ (1 +np.pi), ))
n1 = n - n0

xs = np.concatenate((
        np.stack([np.linspace(0, -r , n0), np.full((n0,), -1.e-8)]),
        r * np.stack([np.cos(np.linspace(-np.pi, np.pi , n1)),np.sin(np.linspace(-np.pi, np.pi , n1))]), 
        np.stack([np.linspace(-r, 0 , n0), np.full((n0,), 1.e-8)]),
    ), axis =1)

xs_def = xs + disp_scale * disp(xs)

fig, ax = plt.subplots(figsize=(10.5,5), ncols = 2)
ax[0].plot(xs[0], xs[1], "b-", label="non-deformed");
ax[1].plot(xs_def[0], xs_def[1], "r-.", label="deformed");

"""
Explanation: Asymptotic Displacement Field of Crack from Linear Elasticity
End of explanation
"""


sim = md.atoms.import_cfg("configs/Fe_300K.cfg");
a = sim.H[0][0]

"""
Explanation: Configuration
Create a $[\bar{1}10]\times\frac{1}{2}[111]\times[11\bar{2}]$ cell
start with a $[100]\times[010]\times[001]$ cell
End of explanation
"""


sim.cell_change([[-1,1,0],[1,1,1],[1,1,-2]])

"""
Explanation: Create a $[\bar{1}10]\times[111]\times[11\bar{2}]$ cell
End of explanation
"""


H = np.array(sim.H);
def _(x):
    if x[1] > 0.5*H[1, 1] - 1.0e-8:
        return False;
    else:
        x[1] *= 2.0;
sim.do(_);

_ = np.full((3,3), 0.0)
_[1, 1] = -0.5
sim.strain(_)

"""
Explanation: Remove half of the atoms and readjust the position of remaining
Now one needs to cut the cell in half in $[111]$ direction. We can achive this in three steps:

Remove the atoms that are above located above $\frac{1}{2}[111]$
Double the position of the remiaing atoms in the said direction
Shrink the box affinly to half on that direction
End of explanation
"""


H = np.array(sim.H);
displace(sim,np.array([sim.H[0][0]/6.0, sim.H[1][1]/6.0, sim.H[2][2]/6.0]))

"""
Explanation: Readjust the postions
End of explanation
"""


max_natms=100000
H=np.array(sim.H);
n_per_area=sim.natms/(H[0,0] * H[1,1]);
_ =np.sqrt(max_natms/n_per_area);
N0 = np.array([
    np.around(_ / sim.H[0][0]),
    np.around(_ / sim.H[1][1]), 
    1], dtype=np.int32)

# make sure in 1 direction it is an even number
if N0[1] % 2 == 1:
    N0[1] += 1

sim *= N0;

"""
Explanation: Replicating the unit cell
End of explanation
"""


vaccum = 100.0

H = np.array(sim.H);
H_new = np.array(sim.H);
H_new[0][0] += vaccum
H_new[1][1] += vaccum
resize(sim, H_new, H.sum(axis=0) * 0.5)

"""
Explanation: Add vacuum
End of explanation
"""


_ = np.array([[-1,1,0],[1,1,1],[1,1,-2]], dtype=np.float);
Q = np.linalg.inv(np.sqrt(_ @ _.T)) @ _;
C = rot(cubic(1.3967587463636366,0.787341583191591,0.609615090769241),Q)
disp = crack(C)

"""
Explanation: Get the displacement field for this configuration
End of explanation
"""


fixed_layer_thickness = 20.0
intensity = 0.5
rate = 0.001

H = np.array(sim.H);
ctr = H.sum(axis=0) * 0.5

lim = np.array([H[0, 0], H[1, 1]]) 
lim -= vaccum;
lim *= 0.5
lim -= fixed_layer_thickness

def _(x, x_d, x_dof):
    
    x_rel = x[:2] - ctr[:2]
    u = disp(x_rel)
    x[:2] += intensity * u
    
    if (np.abs(x_rel) < lim).sum() != 2:
        x_d[:2] = rate * u
        x_dof[0] = False;
        x_dof[1] = False;
    
sim.do(_)

md.export_cfg("", extra_vecs=["x_dof"] )(sim, "dumps/crack.cfg")

"""
Explanation: Impose the diplacement field and other boundary conditions
End of explanation
"""


sim.kB = 8.617330350e-5
sim.hP = 4.13566766225 * 0.1 * np.sqrt(1.60217656535/1.66053904020)
sim.create_temp(300.0, 846244)

"""
Explanation: assign intial velocities
End of explanation
"""


sim.add_elem('H',1.007940)

"""
Explanation: add hydrogen to the system
End of explanation
"""


# GPa and Kelvin
def mu(p,T):
    return -2.37+0.0011237850013293155*T+0.00004308665175*T*np.log(p)-0.000193889932875*T*np.log(T);

muvt = md.muvt(mu(1.0e-3,300.0), 300.0, 0.1, 'H', 73108204);
muvt.nevery = 100;
muvt.nattempts=40000;
muvt.ntally=1000;
muvt.export=md.export_cfg('dumps/dump',10000)

"""
Explanation: define ensemble
muvt
End of explanation
"""


muvt.run(sim,100000);

"""
Explanation: run gcmc
End of explanation
"""


