"""
Financial Options & Market Prediction Expert Model — T4 Optimized
"""
import os
import sys
import traceback

# Force unbuffered output so logs stream properly
os.environ["PYTHONUNBUFFERED"] = "1"

def log(msg):
    print(msg, flush=True)

try:
    import torch
    from datasets import load_dataset, concatenate_datasets
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import LoraConfig, prepare_model_for_kbit_training
    from trl import SFTTrainer, SFTConfig
    log("✅ All imports successful")
except Exception as e:
    log(f"❌ Import error: {e}")
    sys.exit(1)

# ============================================================================
MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.3"
HUB_MODEL_ID = "Saksham7772/FinOptions-Mistral-7B"
OUTPUT_DIR = "./finopt-mistral-7b-qlora"

SYSTEM_PROMPT = (
    "You are a quantitative financial analyst and options trading expert. "
    "For every analysis you provide:\n"
    "1. Identify which input data features are most influential "
    "(e.g., implied volatility, volume, earnings, macro indicators, sentiment)\n"
    "2. Explain the directional impact of each feature on the prediction "
    "(bullish/bearish/neutral and why)\n"
    "3. Provide your market prediction or options strategy recommendation with clear reasoning\n"
    "4. Express your confidence level and key risk factors\n"
    "Think step by step before answering."
)

# ============================================================================
# DATASET PREPARATION
# ============================================================================
log("=" * 60)
log("STEP 1/5: Loading datasets...")
log("=" * 60)

try:
    ds_sujet = load_dataset("sujet-ai/Sujet-Finance-Instruct-177k", split="train")
    log(f"  Sujet Finance: {len(ds_sujet)} rows")

    def convert_sujet(example):
        system = (example.get("system_prompt") or "").strip()
        if not system:
            system = SYSTEM_PROMPT
        user = (example.get("user_prompt") or "").strip()
        answer = (example.get("answer") or "").strip()
        if not user or not answer:
            return {"messages": None}
        return {"messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
            {"role": "assistant", "content": answer},
        ]}

    ds_sujet = ds_sujet.map(convert_sujet, remove_columns=ds_sujet.column_names, num_proc=4)
    ds_sujet = ds_sujet.filter(lambda x: x["messages"] is not None, num_proc=4)
    log(f"  Sujet after conversion: {len(ds_sujet)} rows")

    ds_alpaca = load_dataset("gbharti/finance-alpaca", split="train")
    log(f"  Finance Alpaca: {len(ds_alpaca)} rows")

    def convert_alpaca(example):
        instruction = (example.get("instruction") or "").strip()
        inp = (example.get("input") or "").strip()
        output = (example.get("output") or "").strip()
        if not instruction or not output:
            return {"messages": None}
        user_content = instruction + (f"\n\n{inp}" if inp else "")
        return {"messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_content},
            {"role": "assistant", "content": output},
        ]}

    ds_alpaca = ds_alpaca.map(convert_alpaca, remove_columns=ds_alpaca.column_names, num_proc=4)
    ds_alpaca = ds_alpaca.filter(lambda x: x["messages"] is not None, num_proc=4)
    log(f"  Alpaca after conversion: {len(ds_alpaca)} rows")

    ds_500k = load_dataset("Josephgflowers/Finance-Instruct-500k", split="train")
    log(f"  Finance 500K: {len(ds_500k)} rows")

    def convert_500k(example):
        system = (example.get("system") or "").strip()
        if not system:
            system = SYSTEM_PROMPT
        user = (example.get("user") or "").strip()
        assistant = (example.get("assistant") or "").strip()
        if not user or not assistant:
            return {"messages": None}
        return {"messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant},
        ]}

    ds_500k = ds_500k.map(convert_500k, remove_columns=ds_500k.column_names, num_proc=4)
    ds_500k = ds_500k.filter(lambda x: x["messages"] is not None, num_proc=4)
    log(f"  500K after conversion: {len(ds_500k)} rows")

    combined = concatenate_datasets([ds_sujet, ds_alpaca, ds_500k])
    combined = combined.shuffle(seed=42)
    split = combined.train_test_split(test_size=0.01, seed=42)
    train_dataset = split["train"]
    eval_dataset = split["test"]
    log(f"✅ COMBINED: {len(combined)} rows | Train: {len(train_dataset)} | Eval: {len(eval_dataset)}")

except Exception as e:
    log(f"❌ Dataset error: {e}")
    traceback.print_exc()
    sys.exit(1)

# ============================================================================
# MODEL
# ============================================================================
log("=" * 60)
log("STEP 2/5: Loading model...")
log("=" * 60)

try:
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.float16,
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        quantization_config=bnb_config,
        device_map="auto",
        torch_dtype=torch.float16,
        trust_remote_code=True,
    )
    model = prepare_model_for_kbit_training(model)
    model.config.use_cache = False

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    log(f"✅ Model loaded: {MODEL_ID} ({model.num_parameters():,} params)")

    # Check GPU memory
    if torch.cuda.is_available():
        mem = torch.cuda.memory_allocated() / 1024**3
        total = torch.cuda.get_device_properties(0).total_memory / 1024**3
        log(f"  GPU memory: {mem:.1f}GB / {total:.1f}GB used")
except Exception as e:
    log(f"❌ Model loading error: {e}")
    traceback.print_exc()
    sys.exit(1)

# ============================================================================
# LoRA + SFT CONFIG
# ============================================================================
log("=" * 60)
log("STEP 3/5: Configuring training...")
log("=" * 60)

try:
    peft_config = LoraConfig(
        r=64, lora_alpha=128,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
    )

    sft_config = SFTConfig(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        per_device_train_batch_size=1,
        per_device_eval_batch_size=1,
        gradient_accumulation_steps=16,
        optim="paged_adamw_8bit",
        learning_rate=2e-4,
        max_grad_norm=0.3,
        weight_decay=0.001,
        warmup_ratio=0.03,
        lr_scheduler_type="cosine",
        # NO fp16, NO bf16 — pure fp32 training to avoid T4 AMP issues
        bf16=False,
        fp16=False,
        max_length=2048,
        packing=False,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        eval_strategy="steps",
        eval_steps=2000,
        save_strategy="steps",
        save_steps=2000,
        save_total_limit=3,
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        disable_tqdm=True,
        logging_strategy="steps",
        logging_steps=10,
        logging_first_step=True,
        report_to="none",
        push_to_hub=True,
        hub_model_id=HUB_MODEL_ID,
        hub_strategy="every_save",
        seed=42,
        dataloader_num_workers=2,
    )
    log("✅ Config created")
except Exception as e:
    log(f"❌ Config error: {e}")
    traceback.print_exc()
    sys.exit(1)

# ============================================================================
# TRAINER
# ============================================================================
log("=" * 60)
log("STEP 4/5: Creating trainer and tokenizing...")
log("=" * 60)

try:
    trainer = SFTTrainer(
        model=model,
        args=sft_config,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        peft_config=peft_config,
        processing_class=tokenizer,
    )
    log(f"✅ Trainer created")
    log(f"  Trainable params: {sum(p.numel() for p in trainer.model.parameters() if p.requires_grad):,}")
    log(f"  Total steps: {trainer.state.max_steps if hasattr(trainer.state, 'max_steps') else 'unknown'}")

    # Check param dtypes
    dtype_counts = {}
    for name, p in trainer.model.named_parameters():
        dt = str(p.dtype)
        dtype_counts[dt] = dtype_counts.get(dt, 0) + 1
    log(f"  Param dtypes: {dtype_counts}")

    if torch.cuda.is_available():
        mem = torch.cuda.memory_allocated() / 1024**3
        total = torch.cuda.get_device_properties(0).total_memory / 1024**3
        log(f"  GPU memory after trainer: {mem:.1f}GB / {total:.1f}GB")

except Exception as e:
    log(f"❌ Trainer creation error: {e}")
    traceback.print_exc()
    sys.exit(1)

# ============================================================================
# TRAIN
# ============================================================================
log("=" * 60)
log("STEP 5/5: TRAINING...")
log(f"  Model: {MODEL_ID}")
log(f"  Dataset: {len(train_dataset)} examples")
log(f"  Epochs: 1 | Batch: 1 | GradAccum: 16 | EffBatch: 16")
log(f"  LoRA r=64 α=128 | LR: 2e-4 | fp32 training")
log(f"  Push to: {HUB_MODEL_ID}")
log("=" * 60)

try:
    trainer.train()
    log("✅ Training complete!")
except Exception as e:
    log(f"❌ Training error: {e}")
    traceback.print_exc()
    # Still try to save whatever we have
    try:
        log("Attempting to save partial model...")
        trainer.save_model(OUTPUT_DIR)
        tokenizer.save_pretrained(OUTPUT_DIR)
        log("Partial model saved")
    except:
        log("Could not save partial model")
    sys.exit(1)

# ============================================================================
# SAVE & PUSH
# ============================================================================
log("Saving final model...")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
log(f"✅ DONE! Model at: https://huggingface.co/{HUB_MODEL_ID}")