# Set up lattice
vector_a = [5.13498, 0.0, 0.0]*Angstrom
vector_b = [0.0, -3.45897, 0.0]*Angstrom
vector_c = [0.0, 0.0, 50.0]*Angstrom
lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
elements = [Oxygen, Oxygen, Carbon, Oxygen, Lithium, Lithium, Oxygen, Oxygen,
            Carbon, Oxygen, Lithium, Lithium, Oxygen, Oxygen, Carbon, Oxygen,
            Lithium, Lithium, Oxygen, Oxygen, Carbon, Oxygen, Lithium, Lithium,
            Oxygen, Oxygen, Lithium, Lithium, Oxygen, Oxygen, Lithium, Lithium,
            Oxygen, Oxygen, Lithium, Lithium, Oxygen, Oxygen, Lithium, Lithium,
            Oxygen, Oxygen, Lithium, Lithium, Oxygen, Oxygen, Lithium, Lithium,
            Oxygen, Oxygen, Lithium, Lithium, Oxygen, Oxygen, Lithium, Lithium]

# Define coordinates
fractional_coordinates = [[ 0.380624212955,  0.            ,  0.189439465318],
                          [ 0.762530782053,  0.            ,  0.212239160329],
                          [ 0.506364858149,  0.            ,  0.212390556065],
                          [ 0.382559159796,  0.            ,  0.235559319456],
                          [ 0.004356755633,  0.            ,  0.243099403578],
                          [ 0.500000107621,  0.500000430093,  0.27312399014 ],
                          [ 0.880624320575,  0.500000430093,  0.281832253384],
                          [ 0.262530674431,  0.500000430093,  0.304631948395],
                          [ 0.006364750527,  0.500000430093,  0.304783344131],
                          [ 0.882559267416,  0.500000430093,  0.327952107522],
                          [ 0.504356863254,  0.500000430093,  0.335492191644],
                          [-0.000000000002,  0.            ,  0.365516778206],
                          [ 0.380624212953,  0.            ,  0.37422504145 ],
                          [ 0.762530782051,  0.            ,  0.397024736461],
                          [ 0.506364858147,  0.            ,  0.397176132197],
                          [ 0.382559159794,  0.            ,  0.420344895588],
                          [ 0.004356755631,  0.            ,  0.42788497971 ],
                          [ 0.500000107619,  0.500000430093,  0.457909566272],
                          [ 0.880624320573,  0.500000430093,  0.466617829516],
                          [ 0.262530674429,  0.500000430093,  0.489417524527],
                          [ 0.006364750525,  0.500000430093,  0.489568920263],
                          [ 0.882559267414,  0.500000430093,  0.512737683654],
                          [ 0.504356863252,  0.500000430093,  0.520277767776],
                          [-0.000000000004,  0.            ,  0.526302354338],
                          [ 0.333331921577, -0.000000000001,  0.549168707379],
                          [ 0.833332029199,  0.500000430093,  0.549168707379],
                          [ 0.666666822231, -0.000000000001,  0.564833629128],
                          [ 0.166666714609,  0.500000430093,  0.564833629128],
                          [ 0.333331921577, -0.000000000001,  0.580498550875],
                          [ 0.833332029199,  0.500000430093,  0.580498550875],
                          [ 0.            ,  0.            ,  0.603364903917],
                          [ 0.500000107622,  0.500000430093,  0.603364903917],
                          [ 0.666668231059, -0.000000000001,  0.626145804208],
                          [ 0.166668123438,  0.500000430093,  0.626145804208],
                          [ 0.333333375715, -0.000000000001,  0.641898174213],
                          [ 0.833333483337,  0.500000430093,  0.641898174213],
                          [ 0.666668231059, -0.000000000001,  0.657650544219],
                          [ 0.166668123438,  0.500000430093,  0.657650544219],
                          [ 0.            ,  0.            ,  0.68043144451 ],
                          [ 0.500000107622,  0.500000430093,  0.68043144451 ],
                          [ 0.333331921577, -0.000000000001,  0.703297797551],
                          [ 0.833332029199,  0.500000430093,  0.703297797551],
                          [ 0.666666822231, -0.000000000001,  0.7189627193  ],
                          [ 0.166666714609,  0.500000430093,  0.7189627193  ],
                          [ 0.333331921577, -0.000000000001,  0.734627641047],
                          [ 0.833332029199,  0.500000430093,  0.734627641047],
                          [ 0.            ,  0.            ,  0.757493994089],
                          [ 0.500000107622,  0.500000430093,  0.757493994089],
                          [ 0.666668231059, -0.000000000001,  0.78027489438 ],
                          [ 0.166668123438,  0.500000430093,  0.78027489438 ],
                          [ 0.333333375715, -0.000000000001,  0.796027264385],
                          [ 0.833333483337,  0.500000430093,  0.796027264385],
                          [ 0.666668231059, -0.000000000001,  0.811779634391],
                          [ 0.166668123438,  0.500000430093,  0.811779634391],
                          [ 0.            ,  0.            ,  0.834560534682],
                          [ 0.500000107622,  0.500000430093,  0.834560534682]]

# Set up configuration
bulk_configuration = BulkConfiguration(
    bravais_lattice=lattice,
    elements=elements,
    fractional_coordinates=fractional_coordinates
    )

# Add tags
bulk_configuration.addTags('Left Interface',   [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12,
                                                13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23])
bulk_configuration.addTags('Right Interface',  [24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36,
                                                37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49,
                                                50, 51, 52, 53, 54, 55])
bulk_configuration.addTags('li2co3_electrode', [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
bulk_configuration.addTags('li2o2_electrode',  [40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52,
                                                53, 54, 55])

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
#----------------------------------------
# Basis Set
#----------------------------------------
basis_set = [
    GGABasis.Lithium_DoubleZetaPolarized,
    GGABasis.Carbon_DoubleZetaPolarized,
    GGABasis.Oxygen_DoubleZetaPolarized,
    ]

#----------------------------------------
# Exchange-Correlation
#----------------------------------------
exchange_correlation = GGA.RPBE

numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=(9, 15, 1),
    )

calculator = LCAOCalculator(
    basis_set=basis_set,
    exchange_correlation=exchange_correlation,
    numerical_accuracy_parameters=numerical_accuracy_parameters,
    )

bulk_configuration.setCalculator(calculator)
nlprint(bulk_configuration)
bulk_configuration.update()
nlsave('interface_bulk_opt.nc', bulk_configuration)

# -------------------------------------------------------------
# Optimize Geometry
# -------------------------------------------------------------
indices_0 = [40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52,
             53, 54, 55]
constraints = [RigidBody(indices_0), 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

bulk_configuration = OptimizeGeometry(
        bulk_configuration,
        max_forces=0.05*eV/Ang,
        max_steps=200,
        max_step_length=0.2*Ang,
        constraints=constraints,
        trajectory_filename='interface_bulk_opt.nc',
        disable_stress=True,
        optimizer_method=LBFGS(),
        )
nlsave('interface_bulk_opt.nc', bulk_configuration)
nlprint(bulk_configuration)
