import textwrap

#-------------------------------------------------------------------------------
# make_eq_latconst: Given the pressure and initial lattice constant, makes input
#                   for the lattice constant at the given pressure
#-------------------------------------------------------------------------------
def compute_eq_latconst( Species, ModelName, init_LatConst, Pressure, stack_data_flnm ):

   Text = """
   # Initialise the problem
   units metal
   dimension 3
   boundary p p	p
   atom_style atomic

   # Atom definition
   lattice fcc %s
   region whole block 0 1 0 1 0 1
   create_box 1 whole
   create_atoms	1 region whole
   mass * 1.0

   # Define the interatomic potentials
   pair_style kim LAMMPSvirial %s
   pair_coeff * * %s
   mass * 1.0

   # Settings
   thermo 10000
   thermo_style custom step lx ly lz press pxx pyy pzz pe temp

   # NPT Settings
   reset_timestep 0
   timestep 0.001
   fix 1 all npt temp 0.0001 0.0001 0.005 iso %s %s 100  drag 10.0
   run 200000

   # Print the equilibrated lattice constant
   variable a	equal "lx"
   print "${a}" file %s screen yes

   # Simulation Done
   print "Computed equilibrium Lattice Constant." """ % (str(init_LatConst), ModelName, Species, str(Pressure), str(Pressure), stack_data_flnm)
   return textwrap.dedent(Text)

#-------------------------------------------------------------------------------
# setup_problem: Sets up the problem and defines geometry etc.,
#-------------------------------------------------------------------------------
def setup_problem( Species, ModelName, N_Layers, LatConst, Pressure, Rigid_Grp_SIdx, Rigid_Grp_EIdx, N_Twin_Layers):

   Text = """
   # Initialize the problem
   units metal
   dimension 3
   boundary p p p
   atom_style atomic
   variable latparam equal %s

   # Declare the dimensions of the model in lattice cells
   variable nlat_x equal 1
   variable nlat_y equal 1
   variable nlat_z equal %s

   # Define the box/various dimensions
   variable xdim      equal (${latparam}*sqrt(6)/2)*${nlat_x}
   variable ydim      equal (${latparam}*sqrt(2)/2)*${nlat_y}
   variable zdim      equal (${latparam}/sqrt(3))*${nlat_z}
   variable layer_z   equal (${latparam}/sqrt(3))

   # Area of the stacking fault plane in Angstrom^2
   variable Area equal 2*${xdim}*${ydim}

   # Geometry definition
   lattice fcc ${latparam}

   # Define regions for stacking fault test
   # NOTE: -0.001 in the x- and y- directions only to overcome numerical precision issues

   # Define a region for each  of the top N_layers/2
   variable ntwin_layers equal ${nlat_z}/2

   variable i loop ${ntwin_layers}
   label start_of_loop_region

      variable j equal ${i}+${ntwin_layers}
      variable zmin equal (${j}-1.0)*${layer_z}-0.001
      variable zmax equal ${j}*${layer_z}-0.001
      region  ${i}  block  -0.001 ${xdim} -0.001 ${ydim}  ${zmin} ${zmax} units box
      next i

   jump SELF start_of_loop_region

   variable i delete
   variable j delete

   # Declare group for the rigid blocks for intrinsic stacking fault
   variable zmin equal (%s-1.0)*${layer_z}-0.001
   variable zmax equal (%s-1.0)*${layer_z}+0.001
   region  stack_region block  -0.001 ${xdim} -0.001 ${ydim}  ${zmin} ${zmax} units box

   # Declare group for the rigid blocks for the extrisic stacking fault
   variable zmin equal (%s-1.0)*${layer_z}-0.001
   variable zmax equal (%s-1.0)*${layer_z}+0.001
   region  twin_region block  -0.001 ${xdim} -0.001 ${ydim}  ${zmin} ${zmax} units box

   # Create simulation box and atoms
   region       whole block  0-0.001  ${xdim}+0.01    0-0.01 ${ydim}+0.01    0-0.01  ${zdim}+0.01   units box
   create_box   1     whole
   lattice fcc ${latparam} orient x 1 1 2 orient y -1 1 0 orient z -1 -1 1
   create_atoms 1 region whole

   # Define the interatomic potentials
   pair_style kim LAMMPSvirial %s
   pair_coeff * * %s
   mass * 1.0

   # Rigid holding settings for each layer as a group
   variable i loop ${ntwin_layers}
   label start_of_loop_group

      group ${i} region ${i}
      next i

   jump SELF start_of_loop_group
   variable i delete

   # Rigid holding settings for the stacking fault group
   group stack_group region stack_region

   # Rigid holding settings for the twinning fault group
   group twin_group region twin_region

   # Compute initial energy and output
   compute peratom all pe/atom
   compute eatoms all reduce sum c_peratom
   thermo 100
   thermo_style custom step pe c_eatoms press pxx pyy pzz pe temp
   #dump config all atom 1000 dump.Stack

   # Perform initial relaxation
   reset_timestep 0
   fix 1 all box/relax x %s y %s z %s vmax 0.01
   min_style cg
   minimize 1e-25 1e-25 10000 10000

   # Twin top N_Twin_Layers
   variable twin_move equal -(1.0*${latparam}/sqrt(6))
   variable i loop %s
   label start_of_loop_twin

      variable k equal ${i}*${twin_move}
      displace_atoms ${i} move ${k} 0.0 0.0 units box
      next i

   jump SELF start_of_loop_twin
   variable i delete
   variable k delete
   minimize 1e-25 1e-25 100000 100000
   variable E equal "c_eatoms"
   variable Eini equal ${E} """ % (str(LatConst), str(N_Layers), str(Rigid_Grp_SIdx), str(Rigid_Grp_EIdx), str(Rigid_Grp_SIdx+1), str(Rigid_Grp_EIdx-1), ModelName, Species, str(Pressure), str(Pressure), str(Pressure), str(N_Twin_Layers)  )

   return  textwrap.dedent(Text)

#-------------------------------------------------------------------------------
# make_stack_twin_test:  Makes LAMMPS input for moving both the stacking planes
# completely one after the other
#-------------------------------------------------------------------------------
def make_stack_twin_test(stack_data_flnm):

   Text = """
   variable outfile string "%s"
   variable n_incr equal 100
   variable inc_x equal ${twin_move}/${n_incr}

   print "Npoints =  1  ${n_incr}   ${n_incr} " file ${outfile} screen no

   # Inital point is zero energy and zero partial fraction
   print "0 0" append ${outfile} screen no

   # Apply the incremental displacement until first stacking fault nucleates
   variable i loop ${n_incr}
   label start_of_loop_1

      displace_atoms stack_group move ${inc_x} 0.0 0.0 units box
      fix 2 all setforce 0 0 NULL
      velocity all zero linear
      minimize 1e-25 1e-25 10000 10000

      variable Ecur equal "c_eatoms"
      variable SFED equal (${Ecur}-${Eini})/${Area}
      variable totdisp equal ${i}/${n_incr}
      print "${totdisp} ${SFED}" append ${outfile} screen no

      next i

   jump SELF start_of_loop_1
   variable i delete

   # Apply the incremental displacement until twinning fault nucleates
   variable i loop ${n_incr}
   label start_of_loop_2

      displace_atoms twin_group move ${inc_x} 0.0 0.0 units box
      fix 2 all setforce 0 0 NULL
      velocity all zero linear
      minimize 1e-25 1e-25 10000 10000

      variable Ecur equal "c_eatoms"
      variable SFED equal (${Ecur}-${Eini})/${Area}
      variable totdisp equal (1.0+${i}/${n_incr})
      print "${totdisp} ${SFED}" append ${outfile} screen no

      next i

   jump SELF start_of_loop_2
   variable i delete

   # SIMULATION DONE
   print "All done" """ % ( stack_data_flnm)

   return  textwrap.dedent(Text)

#-------------------------------------------------------------------------------
# make_refine_us: Makes LAMMPS input for refining gamma_us
#-------------------------------------------------------------------------------
def make_refine_us(SFrac_us, dFrac_us, stack_data_flnm):

   Text = """
   # Create the output file and write header
   variable outfile string "%s"

   print "Refined Values" file ${outfile} screen no

   # Apply the displacement until lower bound of unstable stacking stacking fault
   variable disp equal ${twin_move}*%f
   displace_atoms stack_group move ${disp} 0.0 0.0 units box

   # Relax/minimize in z direction
   fix 2 all setforce 0 0 NULL
   minimize 1e-25 1e-25 10000 10000

   variable Ecur equal "c_eatoms"
   variable SFEDcur equal (${Ecur}-${Eini})/${Area}
   variable SFEDprev equal ${SFEDcur}
   variable Refdisp_us equal ${disp}

   # Refinement of the bottom stack plane motion
   variable d_disp equal 0.5*${twin_move}*%f
   variable i loop 10000
   label start_of_loop_refine_stack

      displace_atoms stack_group move ${d_disp} 0.0 0.0 units box
      variable tmp equal ${Refdisp_us}
      variable Refdisp_us equal ${tmp}+${d_disp}
      run 0

      # Relax in z direction
      min_style cg
      minimize 1e-25 1e-25 10000 10000

      variable Ecur equal "c_eatoms"
      variable SFEDcur equal (${Ecur}-${Eini})/${Area}
      variable tmp_d_disp equal ${d_disp}
      if "(${SFEDcur} < ${SFEDprev})" then "variable d_disp equal -0.5*${tmp_d_disp}"
      variable diffpe equal abs(${SFEDcur}-${SFEDprev})
      if "( ${diffpe} < 1e-10 )" then "jump SELF break_us"

      variable SFEDprev equal ${SFEDcur}
      next i

   jump SELF start_of_loop_refine_stack
   label break_us
   variable i delete
   variable Frac_us equal ${Refdisp_us}/${twin_move}
   print "${Frac_us}  ${SFEDcur}" append "%s" screen no
   """ % (stack_data_flnm, SFrac_us, dFrac_us, stack_data_flnm)

   return textwrap.dedent(Text)

#-------------------------------------------------------------------------------
# make_refine_ut: Makes LAMMPS input for refining gamma_ut
#-------------------------------------------------------------------------------
def make_refine_ut(SFrac_ut, dFrac_ut, stack_data_flnm):

   Text = """
   # Create the output file and write header
   variable outfile string "%s"

   print "Refined Values" file ${outfile} screen no

   # Apply the complete displacement stack block
   variable disp equal ${twin_move}
   displace_atoms stack_group move ${disp} 0.0 0.0 units box

   # Relax in z direction
   fix 2 all setforce 0 0 NULL
   minimize 1e-25 1e-25 10000 10000

   # Apply the displacement to the twin block untill the lower bound of the unstable twin
   variable disp equal ${twin_move}*(%f-1.0)
   displace_atoms twin_group move ${disp} 0.0 0.0 units box

   # Relax in z direction
   minimize 1e-25 1e-25 10000 10000

   variable Ecur equal "c_eatoms"
   variable SFEDcur equal (${Ecur}-${Eini})/${Area}
   variable SFEDprev equal ${SFEDcur}
   variable Refdisp_ut equal ${twin_move}+${disp}

   # Refinement of the twin block motion
   variable d_disp equal 0.5*${twin_move}*%f
   variable i loop 10000
   label start_of_loop_refine_twin

      displace_atoms twin_group move ${d_disp} 0.0 0.0 units box
      variable tmp equal ${Refdisp_ut}
      variable Refdisp_ut equal ${tmp}+${d_disp}
      run 0

      # Relax in z direction
      minimize 1e-25 1e-25 10000 10000

      variable Ecur equal "c_eatoms"
      variable SFEDcur equal (${Ecur}-${Eini})/${Area}
      variable tmp_d_disp equal ${d_disp}
      if "(${SFEDcur} < ${SFEDprev})" then "variable d_disp equal -0.5*${tmp_d_disp}"
      variable diffpe equal abs(${SFEDcur}-${SFEDprev})
      if "( ${diffpe} < 1e-10 )" then "jump SELF break_ut"

      variable SFEDprev equal ${SFEDcur}
      next i

   jump SELF start_of_loop_refine_twin
   label break_ut
   variable i delete
   variable Frac_ut equal ${Refdisp_ut}/${twin_move}
   print "${Frac_ut}  ${SFEDcur}" append "%s" screen no
   """ % (stack_data_flnm, SFrac_ut, dFrac_ut, stack_data_flnm)
   return textwrap.dedent(Text)

#-------------------------------------------------------------------------------
# make_gammasurface_moves: Makes LAMMPS input for moving the stacking block
# and sweep the entire gamma surface
#-------------------------------------------------------------------------------
def make_gammasurface_moves(stack_data_flnm, NxPoints, NyPoints):

   Text = """
   # Create the output file and write header
   variable outfile string "%s"
   variable nx_points equal %s
   variable ny_points equal %s
   variable n_incr_x equal ${nx_points}-1
   variable n_incr_y equal ${ny_points}-1

   variable inc_x equal (-1.0*${latparam}*sqrt(6)/2)/${n_incr_x}
   variable inc_y equal (1.0*${latparam}*sqrt(2)/2)/${n_incr_y}

   print "Header: Printing Gamma Surface" file ${outfile} screen no

   #dump config all atom 1000 dump.Stack
   fix 2 all setforce 0 0 NULL

   # Outer loop: Apply the y - displacement
   variable j loop ${ny_points}
   label start_of_loop_y

      # Apply y-displacement from 2nd step
      if "$j > 1" then &
         "displace_atoms stack_group move 0.0 ${inc_y} 0.0 units box"

      # Inner loop: Apply the x- displacement
      variable i loop ${nx_points}
      label start_of_loop_x

         # Apply x-displacement from 2nd step
         if "$i > 1" then &
            "displace_atoms stack_group move ${inc_x} 0.0 0.0 units box" &

         # Relax in z direction
         minimize 1e-25 1e-25 10000 10000

         variable Ecur equal "c_eatoms"
         variable SFED equal (${Ecur}-${Eini})/${Area}
         variable totdisp_x equal (${i}-1)/${n_incr_x}
         variable totdisp_y equal (${j}-1)/${n_incr_y}
         print "${totdisp_x} ${totdisp_y} ${SFED}" append ${outfile} screen no

         next i

      jump SELF start_of_loop_x
      next j

   jump SELF start_of_loop_y
   variable i delete
   variable j delete

   # SIMULATION DONE
   print "All done" """ % (stack_data_flnm, str(NxPoints), str(NyPoints))
   return textwrap.dedent(Text)