from QATK.Calculators.DFT import *
from QATK.Core import *

# Set up lattice
vector_a = [10.0, 0.0, 0.0]*Angstrom
vector_b = [0.0, 10.0, 0.0]*Angstrom
vector_c = [0.0, 0.0, 18.7]*Angstrom
lattice = UnitCell(vector_a, vector_b, vector_c)

# Define elements
elements = [Silicon, Silicon, Silicon, Silicon, Silicon, Silicon, Silicon,
            Silicon, Silicon, Silicon, Silicon, Silicon, Silicon, Silicon,
            Silicon, Silicon, Silicon, Silicon, Silicon, Silicon, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Hafnium,
            Hafnium, Hafnium, Hafnium, Hafnium, Hafnium, Hafnium, Hafnium,
            Hafnium, Hafnium, Hafnium, Hafnium, Hafnium, Hafnium, Hafnium,
            Hafnium, Hafnium, Hafnium, Hafnium, Hafnium, Hafnium, Hafnium,
            Hafnium, Hafnium, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen, Oxygen,
            Oxygen, Oxygen]

# Define coordinates
fractional_coordinates = [[ 0.170550858449,  0.843015212368,  0.687325541537],
                          [ 0.687098889753,  0.90495219017 ,  0.36699793103 ],
                          [ 0.769365581326,  0.185568429965,  0.301924996764],
                          [ 0.140033391285,  0.748118981264,  0.458277872422],
                          [ 0.372848348185,  0.593906622141,  0.680802540498],
                          [ 0.734571218748,  0.047048691418,  0.644802348181],
                          [ 0.794518091806,  0.4376074087  ,  0.47471497503 ],
                          [ 0.499201658032,  0.309004101203,  0.307472289076],
                          [ 0.640425786001,  0.771600203018,  0.652907172835],
                          [ 0.059665878852,  0.268466409928,  0.333044628015],
                          [ 0.008707916548,  0.966704584655,  0.559566420236],
                          [ 0.671166043686,  0.727736474091,  0.489713941963],
                          [ 0.259450712531,  0.49588778371 ,  0.537797660204],
                          [ 0.051056106687,  0.259600877661,  0.497300362476],
                          [ 0.985805547148,  0.53968756682 ,  0.353780770261],
                          [ 0.869496724192,  0.746996109305,  0.744758105328],
                          [ 0.454647721064,  0.084442696231,  0.559498852508],
                          [ 0.485515894146,  0.34940101078 ,  0.467954200338],
                          [ 0.304466419956,  0.124823848213,  0.715310994033],
                          [ 0.432412941876,  0.85125661244 ,  0.440410332072],
                          [ 0.159022866663,  0.196031280228,  0.749611000167],
                          [ 0.273803785392,  0.813741538274,  0.42089741557 ],
                          [ 0.250747680164,  0.706592766435,  0.698080533227],
                          [ 0.054883382153,  0.66607969835 ,  0.397320258023],
                          [ 0.613629055189,  0.092519954706,  0.58437373626 ],
                          [ 0.660942872409,  0.313148776672,  0.31224823677 ],
                          [ 0.091017215672,  0.212535733358,  0.415298255657],
                          [ 0.649820932257,  0.357306350836,  0.474619178158],
                          [ 0.902445918492,  0.332952901219,  0.514021565623],
                          [ 0.247123465922,  0.940584875735,  0.749725736049],
                          [ 0.796955498957,  0.143457829671,  0.700789455049],
                          [ 0.863029723168,  0.472185760071,  0.399400880194],
                          [ 0.059353119688,  0.121975022895,  0.546172388622],
                          [ 0.519291882407,  0.926047683439,  0.377861934392],
                          [ 0.134657679382,  0.910687589048,  0.614844726222],
                          [ 0.029743891405,  0.799160907822,  0.725134638461],
                          [ 0.418719087578,  0.948076874257,  0.514668846416],
                          [ 0.892575511882,  0.271763782144,  0.329845513803],
                          [ 0.414896669909,  0.499539439842,  0.742404389416],
                          [ 0.758036423958,  0.588804226804,  0.507459678959],
                          [ 0.405612763875,  0.481189659508,  0.499127840886],
                          [ 0.300463261805,  0.498541904027,  0.62342584797 ],
                          [ 0.857505871221,  0.963493227521,  0.594658794977],
                          [ 0.635360455702,  0.789873148537,  0.56876093396 ],
                          [ 0.786250228754,  0.697303808946,  0.672318128652],
                          [ 0.185607565332,  0.639727045022,  0.522257901677],
                          [ 0.516609293772,  0.709856726683,  0.458190635311],
                          [ 0.506573007964,  0.693716523731,  0.679329787563],
                          [ 0.036044808365,  0.862739454988,  0.49221013524 ],
                          [ 0.729334981285,  0.069005234176,  0.355958008566],
                          [ 0.641081109203,  0.920983115532,  0.691542850446],
                          [ 0.42660933467 ,  0.216082260699,  0.507782957223],
                          [ 0.16268907595 ,  0.363484857743,  0.531288056997],
                          [ 0.106284566044,  0.428156214256,  0.339029576352],
                          [ 0.343002674628,  0.100331991005,  0.628854112719],
                          [ 0.434385955892,  0.347817313461,  0.383955303406],
                          [ 0.903454460848,  0.593793349278,  0.279573234428],
                          [ 0.743302566474,  0.854659124436,  0.449193549906],
                          [ 0.43316212074 ,  0.189378471233,  0.26753853989 ],
                          [ 0.768069989669,  0.813862310592,  0.311055895056],
                          [ 0.22515120266 ,  0.408450517702,  0.998501355074],
                          [ 0.577400763618,  0.511706745755,  1.013735239892],
                          [ 0.882384893357,  0.085209259361,  0.837283017093],
                          [ 0.55644682471 ,  0.743681161464,  1.166633521031],
                          [ 0.640907503742,  0.324081690791,  0.844120820328],
                          [ 0.725200418318,  0.910755633881,  0.999599455479],
                          [ 0.47262057842 ,  0.157424645188,  0.997978350599],
                          [ 0.26066685045 ,  0.194080283796,  0.852631782916],
                          [ 1.00055945418 ,  0.44348860499 ,  0.856015020965],
                          [ 0.975368349014,  0.658052937841,  1.002680123615],
                          [ 0.37767599648 ,  0.552628652178,  0.847006370112],
                          [ 0.058933919714,  0.244020745645,  1.155556636823],
                          [ 0.404259503331,  0.335204142636,  1.159752854316],
                          [ 0.329310616479,  0.761754521567,  1.012542931288],
                          [ 0.825128480522,  0.261737487575,  1.011679013423],
                          [ 0.747289282525,  0.674492278222,  0.863895264418],
                          [ 0.157754660406,  0.590970635969,  1.164106989214],
                          [ 0.502436314746,  0.940609682498,  0.851114465044],
                          [ 0.078567905971,  0.011668983505,  1.010630977282],
                          [ 0.808352806465,  0.494534029993,  1.161421859279],
                          [ 0.306402365323,  0.990010193209,  1.166389762349],
                          [ 0.643847092441,  0.086876773235,  1.157248659388],
                          [ 0.135225685557,  0.825005518997,  0.847117228588],
                          [ 0.905074009661,  0.847413946636,  1.168876795669],
                          [ 0.610387337265,  0.545719185931,  1.12468579728 ],
                          [ 0.19108833072 ,  0.383247723368,  0.872385223482],
                          [ 0.621519758571,  0.296848735983,  0.953867586568],
                          [ 0.610591150969,  0.291314882787,  1.178508639172],
                          [ 0.36184413128 ,  0.796543853575,  0.882610620114],
                          [ 0.440478372831,  0.348648990915,  0.839430729847],
                          [ 0.269752461442,  0.954416671097,  1.269877246948],
                          [ 0.520434971418,  0.95119602124 ,  1.206016339332],
                          [ 0.685237973028,  0.867121876314,  0.869289741425],
                          [ 0.349021071588,  0.527273453531,  1.194760514317],
                          [ 0.683403741231,  0.118546455897,  1.046166122738],
                          [ 0.694820282661,  0.884871887051,  1.107793312643],
                          [ 0.277327708101,  0.215394693877,  0.968218089718],
                          [ 0.027514865996,  0.44214624315 ,  1.197760698903],
                          [ 0.682530671469,  0.127093748291,  0.82116961154 ],
                          [ 0.112793571553,  0.798558964365,  0.954724788858],
                          [ 0.184374583874,  0.61943871782 ,  1.049940315566],
                          [ 0.200646830437,  0.376398198211,  1.105395660319],
                          [ 0.439449574046,  0.11943244186 ,  0.882975812917],
                          [ 0.855657421553,  0.293276620516,  1.125277674218],
                          [ 0.943775130449,  0.62949764056 ,  1.110075029347],
                          [ 0.933754701842,  0.889267438929,  0.829285744935],
                          [ 0.363618042281,  0.548263223513,  0.957371634943],
                          [ 0.836299261153,  0.049246062307,  1.201233902708],
                          [ 0.096473585424,  0.782500823225,  1.193601027722],
                          [ 0.862344469027,  0.051253703858,  0.950001378272],
                          [ 0.77902131305 ,  0.461711792039,  1.027454334878],
                          [ 0.518688880628,  0.695603442939,  1.269937533852],
                          [ 0.794192530185,  0.48689573664 ,  1.272816036445],
                          [ 0.359579530736,  0.792936506518,  1.124961977891],
                          [ 0.867234788445,  0.290913325432,  0.885901690833],
                          [ 0.935518080846,  0.872131867822,  1.051656210478],
                          [ 0.109885809204,  0.043230642921,  1.124280714633],
                          [ 0.762223246257,  0.708420656426,  1.19612886137 ],
                          [ 0.067158074337,  0.171447149658,  1.261393949271],
                          [ 0.275816416975,  0.959909012218,  1.025944340966],
                          [ 0.115224044204,  0.047161446407,  0.886333248224],
                          [ 0.598666917593,  0.541581220675,  0.887157286055],
                          [ 0.946802147881,  0.632092052642,  0.87847675274 ],
                          [ 0.780292941943,  0.714954170078,  0.971305044164],
                          [ 0.52786645461 ,  0.710729544729,  1.025962370327],
                          [ 0.027816662012,  0.462435878033,  0.96768550491 ],
                          [ 0.257507888539,  0.212019290014,  1.199343794286],
                          [ 0.024580882451,  0.211752068557,  1.030340726938],
                          [ 0.531133946751,  0.9594357516  ,  0.966241004839],
                          [ 0.188829940459,  0.62878953407 ,  0.825291931102],
                          [ 0.440498559579,  0.126637465648,  1.107235421155],
                          [ 0.436293242241,  0.367742739966,  1.0471675934  ]]

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


# %% Set LCAOCalculator

# %% LCAOCalculator

# ----------------------------------------
# Exchange-Correlation
# ----------------------------------------
exchange_correlation = GGAHalf.PBE

k_point_sampling = KpointDensity(
    density_a=4.0 * Angstrom, density_b=4.0 * Angstrom, density_c=4.0 * Angstrom
)

numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=k_point_sampling,
    density_mesh_cutoff=125 * Hartree,
)

model = PretrainedGridValuesModels.density_HKMG_PBEHalf_PD_Medium_Y2026

grid_values_predictor = RangeSeparatedDensityPredictor(
    model=model,
)

algorithm_parameters = AlgorithmParameters(grid_values_predictor=grid_values_predictor)

calculator = LCAOCalculator(
    exchange_correlation=exchange_correlation,
    numerical_accuracy_parameters=numerical_accuracy_parameters,
    checkpoint_handler=NoCheckpointHandler,
    algorithm_parameters=algorithm_parameters,
)


# %% Set Calculator

bulk_configuration.setCalculator(calculator)

bulk_configuration.update()

nlsave('SiOHf_results.hdf5', bulk_configuration)

