"""Flow-CoPD configs (single A100 80GB). Loaded via:
  accelerate launch flow_copd/train_sd3_copd.py --config flow_copd/config_copd.py:text_a2_copd
"""
import os
import imp
import ml_collections

# load flow_grpo's base.py schema by absolute path (we live outside that tree)
_FLOW_GRPO = os.path.join(os.path.dirname(__file__), "..", "third_party", "flow_grpo")
_FLOW_GRPO = os.path.abspath(_FLOW_GRPO)
base = imp.load_source("base", os.path.join(_FLOW_GRPO, "config", "base.py"))

_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
TEACHERS = {
    "text":      os.path.join(_ROOT, "checkpoints/teachers/text"),       # T2: text rendering / OCR
    "pickscore": os.path.join(_ROOT, "checkpoints/teachers/pickscore"),  # T1: human preference
}
DATASET = os.path.join(_FLOW_GRPO, "dataset")


def get_config(name):
    """ml_collections config_flags dispatcher: `config_copd.py:NAME` -> NAME()."""
    return globals()[name]()


def copd_base():
    """Single-A100 OPD base: SD3.5-M + LoRA, 512 res, 10 denoising steps."""
    config = base.get_config()
    config.pretrained.model = "stabilityai/stable-diffusion-3.5-medium"
    config.use_lora = True
    config.resolution = 512
    config.mixed_precision = "fp16"
    config.activation_checkpointing = True          # single-GPU memory

    # --- sampling (on-policy SDE rollouts; small group to fit 1 GPU) ---
    config.sample.num_steps = 10
    config.sample.eval_num_steps = 40
    config.sample.guidance_scale = 4.5
    config.sample.train_batch_size = 4
    config.sample.num_image_per_prompt = 4          # group size G for advantage
    config.sample.num_batches_per_epoch = 2
    config.sample.test_batch_size = 8
    config.sample.global_std = True
    config.sample.noise_level = 0.7

    # --- training ---
    # train.batch_size MUST equal sample.train_batch_size: the per-step training
    # micro-batch = total//num_batches_per_epoch = sample.train_batch_size, and the
    # CFG-teacher path slices train_neg_prompt_embeds[:micro] (sized by train.batch_size).
    config.train.batch_size = 4
    config.train.gradient_accumulation_steps = 1
    config.train.num_inner_epochs = 1
    config.train.learning_rate = 3e-4
    config.train.timestep_fraction = 0.99
    config.train.cfg = False                        # no-CFG OPD (RL = CFG distillation; saves a forward)
    config.train.ema = True
    config.train.beta = 0.0                         # no GRPO-KL; OPD has its own teacher target

    # --- Flow-CoPD specific (consumed by train_sd3_copd.py / copd_loss.py) ---
    config.copd = ml_collections.ConfigDict()
    config.copd.mode = "positive"                   # "positive" (A0) | "contrastive" (A2)
    config.copd.teacher_lora_path = TEACHERS["text"]
    config.copd.teacher_adapter_name = "teacher"
    # CFG-composed teacher (review CRITICAL#2). The jieliu teachers are used at
    # inference with guidance 4.5, so we distill the GUIDED teacher into the no-CFG
    # student (= CFG distillation). SANITY must compare 1.0 vs 4.5 and pick whichever
    # gives a sane distillation target.
    config.copd.teacher_guidance_scale = 4.5
    config.copd.weight_scheme = "uniform"           # per-step time weight: "uniform" | "snr"
    config.copd.lambda_neg = 1.0                    # contrastive negative strength
    config.copd.neg_clamp = 1.0                     # cap on per-sample negative MSE (stability)
    config.copd.adv_thresh = 0.0                    # deadband for "bad" trajectory selection

    config.per_prompt_stat_tracking = True          # group-normalized advantages for trajectory selection
    config.num_epochs = 100000
    config.save_freq = 50
    config.eval_freq = 25
    config.eval_at_start = True                      # run eval at epoch 0 (set False for fast sanity)
    return config


def text_sanity():
    """Gate-0 smoke test of the OPD/contrastive PIPELINE on the real SD3.5 model.
    Uses the PickScore teacher+reward (CLIP-based, paddle-free) to decouple from the
    broken PaddleOCR install — Gate-0 only verifies sampling + teacher/student
    velocity + OPD loss + backward run without OOM/crash. The real text/OCR runs
    (text_a0/a2) need paddle fixed first. Skips eval + checkpoint saving."""
    config = copd_base()
    config.copd.mode = "positive"
    config.copd.teacher_lora_path = TEACHERS["pickscore"]
    config.dataset = os.path.join(DATASET, "ocr")       # any text-prompt set (has train/test.txt)
    config.prompt_fn = "general_ocr"
    config.reward_fn = {"pickscore": 1.0}               # paddle-free reward
    config.eval_at_start = False
    config.eval_freq = 100000
    config.save_freq = 100000
    config.sample.num_batches_per_epoch = 2
    config.run_name = "text_sanity"
    config.save_dir = "logs/copd/text_sanity"
    return config


# ----------------------- TEXT teacher (OCR reward) -----------------------
def sanity_contrastive():
    """Gate-0b: exercise the CONTRASTIVE path (copd_loss + per-step v_neg) on GPU,
    eval/save off. Verifies the method-specific A2 code before long runs."""
    config = text_sanity()
    config.copd.mode = "contrastive"
    config.run_name = "sanity_contrastive"
    config.save_dir = "logs/copd/sanity_contrastive"
    return config


def text_a0_positive():
    """A0 — positive-only OPD (Flow-OPD / DiffusionOPD reproduction)."""
    config = copd_base()
    config.copd.mode = "positive"
    config.copd.teacher_lora_path = TEACHERS["text"]
    config.dataset = os.path.join(DATASET, "ocr")
    config.prompt_fn = "general_ocr"
    config.reward_fn = {"ocr": 1.0}                 # for advantage / monitoring (NOT used as a policy gradient)
    config.save_dir = "logs/copd/text_a0_positive"
    config.run_name = "text_a0_positive"
    return config


def text_a2_copd():
    """A2 — Flow-CoPD = positive OPD + contrastive negative (OUR method)."""
    config = text_a0_positive()
    config.copd.mode = "contrastive"
    config.copd.lambda_neg = 1.0
    config.copd.neg_clamp = 1.0
    config.save_dir = "logs/copd/text_a2_copd"
    config.run_name = "text_a2_copd"
    return config


def text_a1_symmetric():
    """A1 — symmetric-negative control (same magnitude, no asymmetry)."""
    config = text_a2_copd()
    config.copd.mode = "contrastive"
    config.copd.adv_thresh = -1e9                   # treat ALL trajectories as "bad" -> symmetric repulsion control
    config.save_dir = "logs/copd/text_a1_symmetric"
    config.run_name = "text_a1_symmetric"
    return config


# ----------------------- PICKSCORE teacher (preference) -----------------------
def pick_a0_positive():
    config = copd_base()
    config.copd.mode = "positive"
    config.copd.teacher_lora_path = TEACHERS["pickscore"]
    config.dataset = os.path.join(DATASET, "pickscore")
    config.prompt_fn = "general_ocr"               # plain text-prompt dataset loader
    # train on PickScore (proxy); aesthetic weight 0 = logged-only GOLD metric.
    # reward-hacking = reward_pickscore up while reward_aesthetic drops.
    config.reward_fn = {"pickscore": 1.0, "aesthetic": 0.0}
    config.eval_at_start = False                    # skip slow epoch-0 eval; read trends from per-epoch train logs
    config.save_dir = "logs/copd/pick_a0_positive"
    config.run_name = "pick_a0_positive"
    return config


def pick_a2_copd():
    config = pick_a0_positive()
    config.copd.mode = "contrastive"
    config.save_dir = "logs/copd/pick_a2_copd"
    config.run_name = "pick_a2_copd"
    return config


def pick_a1_symmetric():
    """A1 — symmetric-negative control (PickScore): repel ALL trajectories equally
    (adv_thresh huge negative => every traj treated as 'bad'), isolating the value
    of the reward-ASYMMETRY in Flow-CoPD."""
    config = pick_a2_copd()
    config.copd.adv_thresh = -1e9
    config.save_dir = "logs/copd/pick_a1_symmetric"
    config.run_name = "pick_a1_symmetric"
    return config