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

# %% SiGe-nanowire

# Set up lattice
lattice = SimpleTetragonal(40.04546202772588*Angstrom, 5.4306*Angstrom)

# Define elements
elements = [Silicon, Silicon, Silicon, Silicon, Silicon, Silicon, Silicon,
            Germanium, Silicon, Silicon, Silicon, Germanium, Silicon, Silicon,
            Silicon, Silicon, Silicon, Silicon, Germanium, Silicon, Silicon,
            Silicon, Silicon, Germanium, Silicon, Silicon, Silicon, Silicon,
            Silicon, Silicon, Silicon, Silicon, Silicon, Silicon, Silicon,
            Silicon, Silicon, Silicon, Silicon, Silicon, Silicon, Germanium,
            Silicon, Germanium, Silicon, Germanium, Silicon, Silicon, Silicon,
            Silicon, Silicon, Silicon, Germanium, Silicon, Silicon, Silicon,
            Silicon, Silicon, Silicon, Silicon, Silicon, Silicon, Silicon,
            Silicon, Silicon, Silicon, Silicon, Germanium, Silicon, Silicon,
            Silicon, Silicon, Silicon, Silicon, Silicon, Silicon, Silicon,
            Silicon, Silicon, Silicon, Silicon, Silicon, Silicon, Silicon,
            Silicon, Silicon, Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen,
            Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen,
            Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen,
            Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen,
            Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen,
            Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen,
            Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen, Hydrogen,
            Hydrogen]

# Define coordinates
fractional_coordinates = [[ 0.275653560091,  0.409191131607,  0.337672903375],
                          [ 0.274973203032,  0.554398025955,  0.376217527054],
                          [ 0.274357426442,  0.480527000879,  0.872141739685],
                          [ 0.276782472399,  0.62484222877 ,  0.906697043971],
                          [ 0.410428980505,  0.275719356433,  0.337674500394],
                          [ 0.344739304643,  0.343400307609,  0.365245355869],
                          [ 0.378859020463,  0.377143636024,  0.113540240776],
                          [ 0.413397476385,  0.413038314737,  0.367240220502],
                          [ 0.379656462367,  0.30938982927 ,  0.606755991044],
                          [ 0.413767239092,  0.343461888584,  0.862612663041],
                          [ 0.30914457464 ,  0.37744486014 ,  0.599667689902],
                          [ 0.343791476862,  0.412788204851,  0.860939097828],
                          [ 0.308385571713,  0.447694726151,  0.120605460679],
                          [ 0.344025395316,  0.482477320682,  0.366024650483],
                          [ 0.378472464581,  0.517530725185,  0.122045296443],
                          [ 0.412476226106,  0.552747348913,  0.367867291487],
                          [ 0.37864333631 ,  0.449113348401,  0.615782281105],
                          [ 0.412366830981,  0.48253351263 ,  0.873096757126],
                          [ 0.308575137097,  0.517562514899,  0.625520490482],
                          [ 0.34469942154 ,  0.552866749875,  0.878136591219],
                          [ 0.309614007932,  0.587165136461,  0.129320829123],
                          [ 0.344534569215,  0.621420437428,  0.378146607523],
                          [ 0.378500850554,  0.655458970998,  0.116955218115],
                          [ 0.413523954829,  0.691396393346,  0.370844808939],
                          [ 0.379007292136,  0.586804333189,  0.62469285332 ],
                          [ 0.413531366707,  0.621203427725,  0.872922224738],
                          [ 0.310036087022,  0.655114233055,  0.633102269315],
                          [ 0.344507270271,  0.690246513133,  0.866471616149],
                          [ 0.374057954134,  0.725135234169,  0.599054451221],
                          [ 0.553142783134,  0.276406551747,  0.384995617613],
                          [ 0.481899812044,  0.276094231308,  0.862606853326],
                          [ 0.448236755313,  0.308990545692,  0.116500168233],
                          [ 0.482929839566,  0.343913477159,  0.370856107772],
                          [ 0.517324837381,  0.378497605383,  0.122130707572],
                          [ 0.552621389895,  0.412390462289,  0.366993598035],
                          [ 0.517523147301,  0.30939948575 ,  0.624093861981],
                          [ 0.552105750015,  0.344308313623,  0.874500528445],
                          [ 0.448789115874,  0.378151407501,  0.622087547333],
                          [ 0.482387294684,  0.412566070417,  0.874339960633],
                          [ 0.448150676376,  0.448331982944,  0.113164811564],
                          [ 0.481923305537,  0.48207829084 ,  0.366333791755],
                          [ 0.517599744035,  0.517776571721,  0.126333999141],
                          [ 0.552481190039,  0.5524597648  ,  0.388462004915],
                          [ 0.517537789732,  0.447704817601,  0.620194120331],
                          [ 0.552964043373,  0.482549144828,  0.8756029807  ],
                          [ 0.447589525357,  0.517658627528,  0.620958763249],
                          [ 0.482293342594,  0.553216771919,  0.875614714237],
                          [ 0.447747363799,  0.587149390076,  0.125167233671],
                          [ 0.482454181573,  0.621189702963,  0.378348679915],
                          [ 0.517044396419,  0.656378000137,  0.137741985578],
                          [ 0.552710253736,  0.690611793086,  0.386642042852],
                          [ 0.516807944125,  0.586702697918,  0.626855475332],
                          [ 0.552108351786,  0.621069768583,  0.886444299567],
                          [ 0.448668828146,  0.655475617312,  0.629888638875],
                          [ 0.48326715262 ,  0.690616442339,  0.878232058541],
                          [ 0.448206975737,  0.725358742889,  0.106464468798],
                          [ 0.519000226277,  0.723533249401,  0.63870269489 ],
                          [ 0.624475919987,  0.276549241345,  0.910014012515],
                          [ 0.586626048426,  0.309609834875,  0.130468000854],
                          [ 0.621299226294,  0.344228864676,  0.378987560994],
                          [ 0.655689089622,  0.378611968863,  0.120670550679],
                          [ 0.69059556357 ,  0.413171931401,  0.36927131989 ],
                          [ 0.655019497734,  0.309730184854,  0.635517719583],
                          [ 0.690201945945,  0.344892154388,  0.864160238591],
                          [ 0.586624931156,  0.378736971805,  0.624000296741],
                          [ 0.621217864658,  0.413063748382,  0.872370144413],
                          [ 0.586613598987,  0.447271303692,  0.120000239318],
                          [ 0.621322838415,  0.482670790424,  0.377950809473],
                          [ 0.656658161484,  0.517941287178,  0.127705607513],
                          [ 0.690847259349,  0.552379383945,  0.385886102726],
                          [ 0.656421506633,  0.44730835223 ,  0.628620780284],
                          [ 0.690799728246,  0.482709443681,  0.876621162422],
                          [ 0.586720452365,  0.518046455208,  0.636310420356],
                          [ 0.6213227157  ,  0.55213851261 ,  0.886767053034],
                          [ 0.587458255204,  0.586210354644,  0.142005904515],
                          [ 0.621805087096,  0.62128859503 ,  0.387550839416],
                          [ 0.655766577205,  0.655375524116,  0.132855580205],
                          [ 0.656012648361,  0.58664544466 ,  0.640039757179],
                          [ 0.690252010195,  0.620765117343,  0.894215663735],
                          [ 0.587593744335,  0.656256400069,  0.63649959579 ],
                          [ 0.621339858879,  0.690165949649,  0.894102310058],
                          [ 0.590461182995,  0.723783292264,  0.163859819449],
                          [ 0.723560422275,  0.37531117649 ,  0.590063093682],
                          [ 0.723710109411,  0.446952979899,  0.115671270968],
                          [ 0.723905984249,  0.518475693875,  0.63769551626 ],
                          [ 0.724003801525,  0.590212586908,  0.164393463363],
                          [ 0.259817908806,  0.386013349182,  0.154330210541],
                          [ 0.246861052481,  0.425290758093,  0.468845911676],
                          [ 0.24971333063 ,  0.536348907011,  0.222658108901],
                          [ 0.254359207069,  0.576897999179,  0.536432925279],
                          [ 0.254135408068,  0.45783569554 ,  0.710496378282],
                          [ 0.248723261118,  0.498323394287,  0.023805655361],
                          [ 0.247442147188,  0.608557365285,  0.784255610017],
                          [ 0.261883077091,  0.648910140505,  0.088264563799],
                          [ 0.426222372415,  0.246311558359,  0.462622412179],
                          [ 0.386363030121,  0.26087703686 ,  0.15562288557 ],
                          [ 0.323249476922,  0.320916495004,  0.211020685799],
                          [ 0.358176120885,  0.286669138423,  0.760697984066],
                          [ 0.286168795523,  0.355471301705,  0.748500098895],
                          [ 0.287097851448,  0.676306002039,  0.479441145368],
                          [ 0.322700255047,  0.712316424997,  0.022335724136],
                          [ 0.349348582911,  0.740436895411,  0.422998332085],
                          [ 0.389504521941,  0.754285236141,  0.730367276036],
                          [ 0.574952263876,  0.256075193146,  0.551506761355],
                          [ 0.534221262345,  0.25090858664 ,  0.240264233926],
                          [ 0.500344118356,  0.250199730492,  0.006899008456],
                          [ 0.459761087827,  0.256280691252,  0.695126984747],
                          [ 0.427279378139,  0.745237419576,  0.930421492291],
                          [ 0.46770410565 ,  0.751253013018,  0.244013346037],
                          [ 0.500030648613,  0.748693750099,  0.490937401436],
                          [ 0.540709119515,  0.744169556715,  0.80347669786 ],
                          [ 0.648451802353,  0.261699018349,  0.092659276343],
                          [ 0.608254906644,  0.247155747889,  0.787960603895],
                          [ 0.676195815345,  0.286655427585,  0.482583090579],
                          [ 0.713193004633,  0.323674719985,  0.017372038772],
                          [ 0.677920201817,  0.677401599139,  0.286406801003],
                          [ 0.712919025199,  0.64229850813 ,  0.740354546964],
                          [ 0.642998454354,  0.712933897595,  0.742165756091],
                          [ 0.574292067322,  0.752940026625,  1.037992562548],
                          [ 0.614393031719,  0.739128546698,  0.344669503172],
                          [ 0.752865640704,  0.39156534316 ,  0.713114208978],
                          [ 0.738439000405,  0.351354108064,  0.407519174985],
                          [ 0.749232918833,  0.465765460692,  0.260946192122],
                          [ 0.743948505839,  0.425225705008,  0.948035806945],
                          [ 0.744007027018,  0.540370791624,  0.805078952411],
                          [ 0.749516257954,  0.499931495424,  0.491272606989],
                          [ 0.738985395681,  0.614240749652,  0.346004559369],
                          [ 0.753267580062,  0.574174432071,  1.038995264966]]

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


# %% Set LCAOCalculator

# %% LCAOCalculator

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
)

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


# %% Set Calculator

sigenanowire.setCalculator(calculator)

sigenanowire.update()

nlsave('SiGe_nanowire.hdf5', sigenanowire)


# %% Bandstructure

bandstructure = Bandstructure(
    configuration=sigenanowire, route=['Z', 'G', 'Z'], points_per_segment=50
)
nlsave('SiGe_nanowire.hdf5', bandstructure, object_id='Bandstructure SCF')


# %% Set LCAOCalculator

# %% LCAOCalculator

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
)

model = PretrainedGridValuesModels.density_SiGe

grid_values_predictor = DensityPredictor(
    model=model,
)

algorithm_parameters = AlgorithmParameters(grid_values_predictor=grid_values_predictor)

calculator_1 = LCAOCalculator(
    numerical_accuracy_parameters=numerical_accuracy_parameters,
    iteration_control_parameters=NonSelfconsistent,
    checkpoint_handler=NoCheckpointHandler,
    algorithm_parameters=algorithm_parameters,
)


# %% Set Calculator

sigenanowire.setCalculator(calculator_1)

sigenanowire.update()

nlsave('SiGe_nanowire.hdf5', sigenanowire)


# %% Bandstructure

bandstructure_1 = Bandstructure(
    configuration=sigenanowire, route=['Z', 'G', 'Z'], points_per_segment=50
)
nlsave('SiGe_nanowire.hdf5', bandstructure_1, object_id='Bandstructure ML non-SCF')
