module PySchema

include("py_type_conversion.jl")
import Braket: OneQubitProperties, TwoQubitProperties, CoherenceTime, FidelityType, Geometry, AnalogHamiltonianSimulationShotResult, DeviceActionProperties, Direction, DeviceActionType, BraketSchemaBase, PersistedJobDataFormat, DeviceDocumentation, TaskMetadata, AnalogHamiltonianSimulationTaskResult, DeviceConnectivity, Rydberg, RydbergGlobal, AnalogHamiltonianSimulationShotMeasurement, QueraAhsParadigmProperties, Lattice, PersistedJobData, Area, ResultType, ResultFormat, AnalogHamiltonianSimulationShotMetadata, braketSchemaHeader, AbstractProgram
import Braket.IR: Setup, PhysicalField

const analog_hamiltonian_simulation_task_result_v1 = PythonCall.pynew()
const device_execution_window = PythonCall.pynew()
const shifting_field = PythonCall.pynew()
const atom_arrangement = PythonCall.pynew()
const analog_hamiltonian_simulation_task_result_v1 = PythonCall.pynew()
const driving_field = PythonCall.pynew()
const device_action_properties = PythonCall.pynew()
const schema_base = PythonCall.pynew()
const persisted_job_data_v1 = PythonCall.pynew()
const time_series = PythonCall.pynew()
const task_metadata_v1 = PythonCall.pynew()
const physical_field = PythonCall.pynew()
const result_type = PythonCall.pynew()
const hamiltonian = PythonCall.pynew()
const shared_models = PythonCall.pynew()
const schema_header = PythonCall.pynew()
const program_v1 = PythonCall.pynew()

function jl_convert(::Type{AbstractProgram}, x::Py)
    T = Braket.lookup_type(pyconvert(braketSchemaHeader, pygetattr(x, "braketSchemaHeader")))
    return PythonCall.pyconvert_return(pyconvert(T, x))
end

function __init__()
    PythonCall.pycopy!(shifting_field, pyimport("braket.ir.ahs.shifting_field"))
    PythonCall.pycopy!(atom_arrangement, pyimport("braket.ir.ahs.atom_arrangement"))
    PythonCall.pycopy!(analog_hamiltonian_simulation_task_result_v1, pyimport("braket.task_result.analog_hamiltonian_simulation_task_result_v1"))
    PythonCall.pycopy!(driving_field, pyimport("braket.ir.ahs.driving_field"))
    PythonCall.pycopy!(device_action_properties, pyimport("braket.device_schema.device_action_properties"))
    PythonCall.pycopy!(schema_base, pyimport("braket.schema_common.schema_base"))
    PythonCall.pycopy!(persisted_job_data_v1, pyimport("braket.jobs_data.persisted_job_data_v1"))
    PythonCall.pycopy!(time_series, pyimport("braket.ir.ahs.time_series"))
    PythonCall.pycopy!(program_v1, pyimport("braket.ir.ahs.program_v1"))
    PythonCall.pycopy!(task_metadata_v1, pyimport("braket.task_result.task_metadata_v1"))
    PythonCall.pycopy!(physical_field, pyimport("braket.ir.ahs.physical_field"))
    PythonCall.pycopy!(hamiltonian, pyimport("braket.ir.ahs.hamiltonian"))
    PythonCall.pycopy!(shared_models, pyimport("braket.ir.jaqcd.shared_models"))
    PythonCall.pycopy!(schema_header, pyimport("braket.schema_common.schema_header"))
    PythonCall.pyconvert_add_rule("braket.ir.ahs.shifting_field:ShiftingField", ShiftingField, jl_convert)
    PythonCall.pyconvert_add_rule("braket.ir.ahs.atom_arrangement:AtomArrangement", AtomArrangement, jl_convert)
    PythonCall.pyconvert_add_rule("braket.task_result.analog_hamiltonian_simulation_task_result_v1:AnalogHamiltonianSimulationShotResult", AnalogHamiltonianSimulationShotResult, jl_convert)
    PythonCall.pyconvert_add_rule("braket.ir.ahs.driving_field:DrivingField", DrivingField, jl_convert)
    PythonCall.pyconvert_add_rule("braket.device_schema.device_action_properties:DeviceActionType", DeviceActionType, jl_convert)
    PythonCall.pyconvert_add_rule("braket.schema_common.schema_base:BraketSchemaBase", BraketSchemaBase, jl_convert)
    PythonCall.pyconvert_add_rule("braket.jobs_data.persisted_job_data_v1:PersistedJobDataFormat", PersistedJobDataFormat, jl_convert)
    PythonCall.pyconvert_add_rule("braket.device_schema.device_service_properties_v1:DeviceDocumentation", DeviceDocumentation, jl_convert)
    PythonCall.pyconvert_add_rule("braket.ir.ahs.time_series:TimeSeries", TimeSeries, jl_convert)
    PythonCall.pyconvert_add_rule("braket.ir.ahs.program_v1:AHSProgram", AHSProgram, jl_convert)
    PythonCall.pyconvert_add_rule("braket.task_result.task_metadata_v1:TaskMetadata", TaskMetadata, jl_convert)
    PythonCall.pyconvert_add_rule("braket.task_result.analog_hamiltonian_simulation_task_result_v1:AnalogHamiltonianSimulationTaskResult", AnalogHamiltonianSimulationTaskResult, jl_convert)
    PythonCall.pyconvert_add_rule("braket.ir.ahs.program_v1:Setup", Setup, jl_convert)
    PythonCall.pyconvert_add_rule("braket.ir.ahs.physical_field:PhysicalField", PhysicalField, jl_convert)
    PythonCall.pyconvert_add_rule("braket.task_result.analog_hamiltonian_simulation_task_result_v1:AnalogHamiltonianSimulationShotMeasurement", AnalogHamiltonianSimulationShotMeasurement, jl_convert)
    PythonCall.pyconvert_add_rule("braket.jobs_data.persisted_job_data_v1:PersistedJobData", PersistedJobData, jl_convert)
    PythonCall.pyconvert_add_rule("braket.ir.ahs.hamiltonian:Hamiltonian", Hamiltonian, jl_convert)
    PythonCall.pyconvert_add_rule("braket.task_result.analog_hamiltonian_simulation_task_result_v1:AnalogHamiltonianSimulationShotMetadata", AnalogHamiltonianSimulationShotMetadata, jl_convert)
    PythonCall.pyconvert_add_rule("braket.schema_common.schema_header:BraketSchemaHeader", braketSchemaHeader, jl_convert)
    PythonCall.pyconvert_add_rule("braket.ir.ahs.program_v1:AHSProgram", AbstractProgram, jl_convert)
    PythonCall.pyconvert_add_rule("braket.ir.ahs.program_v1:Program", AbstractProgram, jl_convert)
end

end #module
