from cheetah.accelerator import (  # noqa: F401
    BPM,
    Aperture,
    Cavity,
    CustomTransferMap,
    Dipole,
    Drift,
    Element,
    HorizontalCorrector,
    Marker,
    Quadrupole,
    RBend,
    Screen,
    Segment,
    Sextupole,
    Solenoid,
    SpaceChargeKick,
    TransverseDeflectingCavity,
    Undulator,
    VerticalCorrector,
)

cheetah_conversion_rules = {
    "Dipole": Dipole,
    "Solenoid": Solenoid,
    "Quadrupole": Quadrupole,
    "Sextupole": Sextupole,
    "Octupole": Drift,
    "Beam_Position_Monitor": Screen,
    "Beam_Arrival_Monitor": Screen,
    "Bunch_Length_Monitor": Screen,
    "Screen": Screen,
    "RCollimator": Aperture,
    "Collimator": Aperture,
    "Monitor": Screen,
    "Wall_Current_Monitor": Drift,
    "Integrated_Current_Transformer": Drift,
    "Faraday_Cup_Monitor": Drift,
    "RFCavity": Cavity,
    "RFDeflectingCavity": TransverseDeflectingCavity,
    "Aperture": Aperture,
    "Shutter": Drift,
    "Valve": Drift,
    "Bellows": Drift,
    "Cleaner": Drift,
    "lscdrift": Drift,
    "Drift": Drift,
    "csrdrift": Drift,
    "Kicker": HorizontalCorrector,
    "Combined_Corrector": HorizontalCorrector,
    "Marker": Screen,
    "Wakefield": Drift,
    "Laser": Drift,
}
