from esp_ppq.core import NetworkFramework, TargetPlatform
from esp_ppq.parser import (
    AscendExporter,
    EspdlExporter,
    ExtensionExporter,
    MNNExporter,
    NativeExporter,
    NativeImporter,
    NCNNExporter,
    NxpExporter,
    OnnxExporter,
    OnnxParser,
    ONNXRUNTIMExporter,
    OpenvinoExporter,
    PPLBackendExporter,
    QNNDSPExporter,
    TengineExporter,
    TensorRTExporter_JSON,
)
from esp_ppq.quantization.quantizer import (
    AscendQuantizer,
    EspdlCHPreInt16Quantizer,
    EspdlCInt16Quantizer,
    EspdlCQuantizer,
    EspdlHPreInt16Quantizer,
    EspdlInt16Quantizer,
    EspdlQuantizer,
    EspdlS3HPreInt16Quantizer,
    EspdlS3Int16Quantizer,
    EspdlS3Quantizer,
    ExtQuantizer,
    FPGAQuantizer,
    GraphCoreQuantizer,
    MetaxChannelwiseQuantizer,
    MetaxTensorwiseQuantizer,
    MNNQuantizer,
    NCNNQuantizer,
    NXP_Quantizer,
    OnnxruntimeQuantizer,
    OpenvinoQuantizer,
    PPL_DSP_Quantizer,
    PPL_DSP_TI_Quantizer,
    PPLCUDAQuantizer,
    RKNN_PerTensorQuantizer,
    TengineQuantizer,
    TensorRTQuantizer,
    TensorRTQuantizer_FP8,
)

__QUANTIZER_COLLECTION__ = {
    TargetPlatform.PPL_DSP_INT8: PPL_DSP_Quantizer,
    TargetPlatform.PPL_DSP_TI_INT8: PPL_DSP_TI_Quantizer,
    TargetPlatform.SNPE_INT8: PPL_DSP_Quantizer,
    TargetPlatform.QNN_DSP_INT8: PPL_DSP_Quantizer,
    TargetPlatform.TRT_INT8: TensorRTQuantizer,
    TargetPlatform.MNN_INT8: MNNQuantizer,
    TargetPlatform.ASC_INT8: AscendQuantizer,
    TargetPlatform.NCNN_INT8: NCNNQuantizer,
    TargetPlatform.NXP_INT8: NXP_Quantizer,
    TargetPlatform.RKNN_INT8: RKNN_PerTensorQuantizer,
    TargetPlatform.METAX_INT8_C: MetaxChannelwiseQuantizer,
    TargetPlatform.METAX_INT8_T: MetaxTensorwiseQuantizer,
    TargetPlatform.PPL_CUDA_INT8: PPLCUDAQuantizer,
    TargetPlatform.EXTENSION: ExtQuantizer,
    TargetPlatform.FPGA_INT8: FPGAQuantizer,
    TargetPlatform.OPENVINO_INT8: OpenvinoQuantizer,
    TargetPlatform.TENGINE_INT8: TengineQuantizer,
    TargetPlatform.GRAPHCORE_FP8: GraphCoreQuantizer,
    TargetPlatform.TRT_FP8: TensorRTQuantizer_FP8,
    TargetPlatform.ONNXRUNTIME: OnnxruntimeQuantizer,
    TargetPlatform.ESPDL_INT8: EspdlQuantizer,
    TargetPlatform.ESPDL_INT16: EspdlInt16Quantizer,
    TargetPlatform.ESPDL_S3_INT8: EspdlS3Quantizer,
    TargetPlatform.ESPDL_S3_INT16: EspdlS3Int16Quantizer,
    TargetPlatform.ESPDL_H_PRE_INT16: EspdlHPreInt16Quantizer,
    TargetPlatform.ESPDL_S3_H_PRE_INT16: EspdlS3HPreInt16Quantizer,
    TargetPlatform.ESPDL_C_INT8: EspdlCQuantizer,
    TargetPlatform.ESPDL_C_INT16: EspdlCInt16Quantizer,
    TargetPlatform.ESPDL_C_H_PRE_INT16: EspdlCHPreInt16Quantizer,
}


__PARSERS__ = {NetworkFramework.ONNX: OnnxParser, NetworkFramework.NATIVE: NativeImporter}


__EXPORTERS__ = {
    # TargetPlatform.PPL_DSP_INT8:  PPLDSPCaffeExporter,
    # TargetPlatform.PPL_DSP_TI_INT8: PPLDSPTICaffeExporter,
    TargetPlatform.QNN_DSP_INT8: QNNDSPExporter,
    TargetPlatform.PPL_CUDA_INT8: PPLBackendExporter,
    # TargetPlatform.SNPE_INT8:     SNPECaffeExporter,
    TargetPlatform.NXP_INT8: NxpExporter,
    TargetPlatform.ONNX: OnnxExporter,
    TargetPlatform.ONNXRUNTIME: ONNXRUNTIMExporter,
    TargetPlatform.OPENVINO_INT8: OpenvinoExporter,
    # TargetPlatform.CAFFE:         CaffeExporter,
    TargetPlatform.NATIVE: NativeExporter,
    TargetPlatform.EXTENSION: ExtensionExporter,
    TargetPlatform.RKNN_INT8: OnnxExporter,
    TargetPlatform.METAX_INT8_C: ONNXRUNTIMExporter,
    TargetPlatform.METAX_INT8_T: ONNXRUNTIMExporter,
    TargetPlatform.TRT_INT8: TensorRTExporter_JSON,
    TargetPlatform.ASC_INT8: AscendExporter,
    TargetPlatform.TRT_FP8: ONNXRUNTIMExporter,
    TargetPlatform.NCNN_INT8: NCNNExporter,
    TargetPlatform.TENGINE_INT8: TengineExporter,
    TargetPlatform.MNN_INT8: MNNExporter,
    TargetPlatform.ESPDL_INT8: EspdlExporter,
    TargetPlatform.ESPDL_INT16: EspdlExporter,
    TargetPlatform.ESPDL_S3_INT8: EspdlExporter,
    TargetPlatform.ESPDL_S3_INT16: EspdlExporter,
    TargetPlatform.ESPDL_H_PRE_INT16: EspdlExporter,
    TargetPlatform.ESPDL_S3_H_PRE_INT16: EspdlExporter,
    TargetPlatform.ESPDL_C_INT8: EspdlExporter,
    TargetPlatform.ESPDL_C_INT16: EspdlExporter,
    TargetPlatform.ESPDL_C_H_PRE_INT16: EspdlExporter,
}


__all__ = ['__QUANTIZER_COLLECTION__', '__PARSERS__', '__EXPORTERS__']
