from ocelot.cpbd.elements import (
    Solenoid,
    SBend,
    Quadrupole,
    Sextupole,
    Octupole,
    Monitor,
    Marker,
    Cavity,
    Hcor,
    Undulator,
    TDCavity,
    Aperture,
    Drift,
)

ocelot_conversion_rules = {
    "Dipole": SBend,
    "Solenoid": Solenoid,
    "Quadrupole": Quadrupole,
    "Sextupole": Sextupole,
    "Octupole": Octupole,
    "Beam_Position_Monitor": Marker,
    "Beam_Arrival_Monitor": Drift,
    "Bunch_Length_Monitor": Drift,
    "Screen": Marker,
    "Rcollimator": Aperture,
    "Collimator": Aperture,
    "Monitor": Monitor,
    "Wall_Current_Monitor": Drift,
    "Integrated_Current_Transformer": Drift,
    "Faraday_Cup_Monitor": Drift,
    "RFDeflectingCavity": TDCavity,
    "RFCavity": Cavity,
    "Aperture": Aperture,
    "Shutter": Drift,
    "Valve": Drift,
    "Bellows": Drift,
    "Cleaner": Drift,
    "lscrift": Drift,
    "csrdrift": Drift,
    "edrift": Drift,
    "Drift": Drift,
    "Kicker": Hcor,
    "Scatter": Marker,
    "APContour": Marker,
    "Center": Marker,
    "FEL_Modulator": Undulator,
    "Wiggler": Undulator,
    "Charge": Marker,
    "Wakefield": Marker,
    "Watch_Point": Monitor,
    "Combined_Corrector": Hcor,
    "Marker": Marker,
    "Laser": Drift,
    "Plasma": Drift,
}
