"""FLUX.2-klein-4B -> OpenVINO FP16 export + NNCF weight-only INT4 quantization.

Environment (pinned, do not change):
  diffusers==0.37.1 transformers==4.57.6 tokenizers==0.22.0
  huggingface-hub==0.35.1 optimum-intel==2.2.0 optimum==2.3.0
  openvino==2026.4.0 nncf==3.4.0 torch pillow psutil

FP16 export CLI equivalent (reference, same as task example pattern):
  optimum-cli export openvino -m black-forest-labs/FLUX.2-klein-4B \
    --task text-to-image --library diffusers --weight-format fp16 \
    ./flux-2-klein-4b-ov-fp16
  # same syntax, only -m and output dir differ.

INT4 quantization:
  OVQuantizer + OVPipelineQuantizationConfig (pipeline variant of OVConfig)
  transformer + text_encoder: OVWeightQuantizationConfig(
    bits=4, sym=False, group_size=128, group_size_fallback="adjust", ratio=1.0)
  rest: default INT8 weight-only (OVWeightQuantizationConfig() defaults bits=8)
  passed as ov_config=... to quantizer.quantize()
"""
import os
import time
import json
from pathlib import Path

FP16_DIR = "./flux-2-klein-4b-ov-fp16"
INT4_DIR = "./flux-2-klein-4b-ov-int4"
MODEL_ID = "black-forest-labs/FLUX.2-klein-4B"


def export_fp16():
    from optimum.intel import OVDiffusionPipeline

    if os.path.exists(FP16_DIR) and os.listdir(FP16_DIR):
        print(f"FP16 model already exists at {FP16_DIR}, skipping export...")
        print("CLI equivalent: optimum-cli export openvino -m "
              f"{MODEL_ID} --task text-to-image --library diffusers "
              f"--weight-format fp16 {FP16_DIR}")
        return FP16_DIR, 0.0

    print(f"Exporting {MODEL_ID} to FP16 OpenVINO...")
    print("CLI equivalent: optimum-cli export openvino -m "
          f"{MODEL_ID} --task text-to-image --library diffusers "
          f"--weight-format fp16 {FP16_DIR}")
    start = time.time()
    pipeline = OVDiffusionPipeline.from_pretrained(
        MODEL_ID,
        export=True,
        compile=False,
        weight_format="fp16",
        token=os.environ.get("HF_TOKEN"),
    )
    pipeline.save_pretrained(FP16_DIR)
    export_time = time.time() - start
    print(f"FP16 export completed in {export_time:.2f}s")
    return FP16_DIR, export_time


def quantize_int4(fp16_dir):
    from optimum.intel import OVDiffusionPipeline
    from optimum.intel.openvino import (
        OVQuantizer,
        OVConfig,
        OVPipelineQuantizationConfig,
        OVWeightQuantizationConfig,
    )

    print(f"Quantizing {fp16_dir} to INT4...")
    start = time.time()

    pipeline = OVDiffusionPipeline.from_pretrained(fp16_dir, compile=False)
    quantizer = OVQuantizer.from_pretrained(pipeline)

    int4_cfg = OVWeightQuantizationConfig(
        bits=4,
        sym=False,
        group_size=128,
        group_size_fallback="adjust",
        ratio=1.0,
    )
    int8_default = OVWeightQuantizationConfig()  # bits=8 default INT8

    # OVPipelineQuantizationConfig is the pipeline variant of OVConfig;
    # passed as ov_config=OVConfig(quantization_config=...) to quantize().
    # transformer/text_encoder INT4, everything else defaults to INT8.
    ov_config = OVConfig(
        quantization_config=OVPipelineQuantizationConfig(
            quantization_configs={
                "transformer": int4_cfg,
                "text_encoder": int4_cfg,
            },
            default_config=int8_default,
        )
    )

    quantizer.quantize(
        save_directory=INT4_DIR,
        ov_config=ov_config,
    )

    quantize_time = time.time() - start
    print(f"INT4 quantization completed in {quantize_time:.2f}s")
    return INT4_DIR, quantize_time


def main():
    print("=" * 60)
    print("FLUX.2-klein-4B OpenVINO INT4 Quantization")
    print("=" * 60)
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    fp16_dir, export_time = export_fp16()
    int4_dir, quantize_time = quantize_int4(fp16_dir)

    data = {
        "model_id": MODEL_ID,
        "fp16_export_dir": fp16_dir,
        "int4_dir": int4_dir,
        "export_time_seconds": export_time,
        "quantize_time_seconds": quantize_time,
        "total_time_seconds": export_time + quantize_time,
        "quant_config": {
            "transformer": "OVWeightQuantizationConfig(bits=4,sym=False,group_size=128,group_size_fallback=adjust,ratio=1.0)",
            "text_encoder": "OVWeightQuantizationConfig(bits=4,sym=False,group_size=128,group_size_fallback=adjust,ratio=1.0)",
            "others": "default INT8 (OVWeightQuantizationConfig bits=8)",
            "ov_config_class": "OVPipelineQuantizationConfig (pipeline variant of OVConfig)",
        },
    }
    with open(os.path.join(OUTPUT_DIR, "benchmark_quantization.json"), "w") as f:
        json.dump(data, f, indent=2)
    print(f"\nDone. FP16: {fp16_dir} INT4: {int4_dir}")


if __name__ == "__main__":
    main()