import os
import argparse
import math
import shutil

import numpy as np
from PIL import Image
import onnx
import onnxsim
import nncase


def parse_model_input_output(model_file, input_shape):
    """
    Parse ONNX model input information.

    Args:
        model_file (str): Path to the ONNX model.
        input_shape (list): Input shape [N, C, H, W].

    Returns:
        onnx.ModelProto: Loaded ONNX model.
        list: List of input tensor dictionaries with keys: name, dtype, shape.
    """
    onnx_model = onnx.load(model_file)

    # Collect all graph inputs
    input_all = [node.name for node in onnx_model.graph.input]
    input_initializer = [node.name for node in onnx_model.graph.initializer]

    # Exclude initializer tensors
    input_names = list(set(input_all) - set(input_initializer))
    input_tensors = [node for node in onnx_model.graph.input if node.name in input_names]

    inputs = []
    for _, e in enumerate(input_tensors):
        onnx_type = e.type.tensor_type
        input_dict = {
            'name': e.name,
            'dtype': onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[onnx_type.elem_type],
            'shape': [
                (dim.dim_value if dim.dim_value != 0 else default_dim)
                for dim, default_dim in zip(onnx_type.shape.dim, input_shape)
            ],
        }
        inputs.append(input_dict)

    return onnx_model, inputs


def onnx_simplify(model_file, dump_dir, input_shape):
    """
    Simplify ONNX model using onnx-simplifier.

    Args:
        model_file (str): Path to the ONNX model.
        dump_dir (str): Directory to save simplified model.
        input_shape (list): Input shape [N, C, H, W].

    Returns:
        str: Path to the simplified ONNX model.
    """
    onnx_model, inputs = parse_model_input_output(model_file, input_shape)

    # Infer shapes first
    onnx_model = onnx.shape_inference.infer_shapes(onnx_model)

    input_shapes = {inp['name']: inp['shape'] for inp in inputs}

    # Simplify ONNX model
    onnx_model, check = onnxsim.simplify(onnx_model, input_shapes=input_shapes)
    assert check, "Simplified ONNX model could not be validated."

    # Save simplified model
    simplified_path = os.path.join(dump_dir, 'simplified.onnx')
    onnx.save_model(onnx_model, simplified_path)
    return simplified_path


def read_model_file(model_file):
    """
    Read model file as bytes.

    Args:
        model_file (str): Path to model file.

    Returns:
        bytes: Model file content.
    """
    with open(model_file, 'rb') as f:
        return f.read()


def generate_data_random(shape, batch):
    """
    Generate random uint8 calibration data.

    Args:
        shape (list): Input shape [N, C, H, W].
        batch (int): Number of samples.

    Returns:
        list: Random calibration data.
    """
    data = []
    for _ in range(batch):
        sample = np.random.randint(0, 256, shape).astype(np.uint8)
        data.append([sample])
    return data


def generate_data(shape, batch, calib_dir):
    """
    Generate calibration data from image dataset.

    Args:
        shape (list): Input shape [N, C, H, W].
        batch (int): Number of calibration samples.
        calib_dir (str): Directory containing calibration images.

    Returns:
        np.ndarray: Calibration dataset.
    """
    img_paths = [os.path.join(calib_dir, p) for p in os.listdir(calib_dir)]
    data = []

    for i in range(batch):
        assert i < len(img_paths), "Calibration images not enough."

        img = Image.open(img_paths[i]).convert('RGB')
        img = img.resize((shape[3], shape[2]), Image.BILINEAR)

        img_data = np.asarray(img, dtype=np.uint8)
        img_data = np.transpose(img_data, (2, 0, 1))  # HWC -> CHW

        data.append([img_data[np.newaxis, ...]])

    return np.array(data)


def main():
    """
    Main entry: ONNX -> nncase compilation -> kmodel generation.
    """
    parser = argparse.ArgumentParser(prog="nncase")
    parser.add_argument("--target", default="k230", type=str, help="Target platform: k230 / cpu")
    parser.add_argument("--model", type=str, required=True, help="ONNX model file")
    parser.add_argument("--dataset", type=str, default="../calibration_pose", help="Calibration dataset directory")
    parser.add_argument("--input_width", type=int, default=320, help="Model input width")
    parser.add_argument("--input_height", type=int, default=320, help="Model input height")
    parser.add_argument("--ptq_option", type=int, default=0, help="PTQ option: 0,1,2,3,4,5")

    args = parser.parse_args()

    # Align input size to multiples of 32
    input_width = int(math.ceil(args.input_width / 32.0)) * 32
    input_height = int(math.ceil(args.input_height / 32.0)) * 32

    # Model input shape (NCHW)
    input_shape = [1, 3, input_height, input_width]

    # Temporary directory
    dump_dir = 'tmp'
    os.makedirs(dump_dir, exist_ok=True)

    # Step 1: Simplify ONNX model
    model_file = onnx_simplify(args.model, dump_dir, input_shape)

    # Step 2: Configure compile options
    compile_options = nncase.CompileOptions()
    compile_options.target = args.target

    # Preprocess settings
    compile_options.preprocess = True
    compile_options.swapRB = False
    compile_options.input_shape = input_shape
    compile_options.input_type = 'uint8'
    compile_options.input_range = [0, 1]
    compile_options.mean = [0, 0, 0]
    compile_options.std = [1, 1, 1]
    compile_options.input_layout = "NCHW"

    # Debug options
    # compile_options.dump_ir = True
    # compile_options.dump_asm = True
    # compile_options.dump_dir = dump_dir
    compile_options.quant_type = 'uint8'

    # Step 3: Create compiler
    compiler = nncase.Compiler(compile_options)

    # Step 4: Import ONNX model
    model_content = read_model_file(model_file)
    import_options = nncase.ImportOptions()
    compiler.import_onnx(model_content, import_options)

    # Step 5: Configure PTQ (Post-Training Quantization)
    ptq_options = nncase.PTQTensorOptions()
    ptq_options.samples_count = 20

    if args.ptq_option == 0:
        ptq_options.calibrate_method = 'NoClip'
        ptq_options.w_quant_type = 'uint8'
        ptq_options.quant_type = 'uint8'
    elif args.ptq_option == 1:
        ptq_options.calibrate_method = 'NoClip'
        ptq_options.w_quant_type = 'int16'
    elif args.ptq_option == 2:
        ptq_options.calibrate_method = 'NoClip'
        ptq_options.quant_type = 'int16'
    elif args.ptq_option == 3:
        ptq_options.calibrate_method = 'Kld'
        ptq_options.w_quant_type = 'uint8'
        ptq_options.quant_type = 'uint8'
    elif args.ptq_option == 4:
        ptq_options.calibrate_method = 'Kld'
        ptq_options.w_quant_type = 'int16'
    elif args.ptq_option == 5:
        ptq_options.calibrate_method = 'Kld'
        ptq_options.quant_type = 'int16'

    # Load calibration data
    calib_data = generate_data(input_shape, ptq_options.samples_count, args.dataset)
    ptq_options.set_tensor_data(calib_data)
    compiler.use_ptq(ptq_options)

    # Step 6: Compile model
    compiler.compile()

    # Step 7: Export kmodel
    kmodel_bytes = compiler.gencode_tobytes()
    base, _ = os.path.splitext(args.model)
    kmodel_name = base + ".kmodel"

    with open(kmodel_name, 'wb') as f:
        f.write(kmodel_bytes)

    # Step 8: Clean temporary files
    if os.path.exists("./tmp"):
        shutil.rmtree("./tmp")
    if os.path.exists("./gmodel_dump_dir"):
        shutil.rmtree("./gmodel_dump_dir")


if __name__ == '__main__':
    main()
