"""
BTC Chart Trading Signal Prediction with Qwen3-VL-4B LoRA.

Usage:
  python predict.py --image chart.png
  python predict.py --adapter ./checkpoint-200 --image chart.png
  python predict.py --adapter LangQuant/LQ-Qwen3-VL-4B-ChartSignal --image chart.png
"""

import json
import argparse

import torch
from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
from peft import PeftModel

BASE_MODEL = "Qwen/Qwen3-VL-4B-Instruct"

SYSTEM_PROMPT = (
    "You are a professional Bitcoin futures trader. "
    "Analyze 15-minute candlestick charts to predict the direction over the next 4 hours."
)
USER_PROMPT = (
    "BTCUSDT 15m chart. Predict the direction for the next 4 hours (16 candles).\n"
    "Respond in JSON."
)


def load_model(adapter_path: str, base_model: str = BASE_MODEL):
    """Load base model + LoRA adapter."""
    print(f"Loading model: {base_model}")
    processor = AutoProcessor.from_pretrained(base_model)
    model = Qwen3VLForConditionalGeneration.from_pretrained(
        base_model,
        torch_dtype=torch.bfloat16,
        device_map="auto",
    )

    print(f"Loading LoRA adapter: {adapter_path}")
    model = PeftModel.from_pretrained(model, adapter_path)
    model.eval()
    return model, processor


def predict(model, processor, image_path: str) -> dict:
    """Run inference on a single chart image."""
    messages = [
        {"role": "system", "content": [
            {"type": "text", "text": SYSTEM_PROMPT},
        ]},
        {"role": "user", "content": [
            {"type": "image", "image": image_path},
            {"type": "text", "text": USER_PROMPT},
        ]},
    ]

    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    with torch.no_grad():
        generated_ids = model.generate(**inputs, max_new_tokens=512)

    trimmed = [
        out_ids[len(in_ids):]
        for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    raw_output = processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]

    # Parse JSON
    text = raw_output.strip()
    if text.startswith("```"):
        text = text.split("\n", 1)[1].rsplit("```", 1)[0].strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return {"raw": raw_output}


def main():
    parser = argparse.ArgumentParser(description="BTC Chart Trading Signal Prediction")
    parser.add_argument("--image", type=str, required=True, help="Path to chart image")
    parser.add_argument("--adapter", type=str, default="LangQuant/LQ-Qwen3-VL-4B-ChartSignal",
                        help="LoRA adapter path or HuggingFace repo ID")
    parser.add_argument("--base-model", type=str, default=BASE_MODEL)
    args = parser.parse_args()

    model, processor = load_model(args.adapter, args.base_model)
    result = predict(model, processor, args.image)

    print("\n" + "=" * 60)
    print("Prediction Result")
    print("=" * 60)
    print(json.dumps(result, indent=2, ensure_ascii=False))

    signal = result.get("signal", "?")
    conf = result.get("confidence", "?")
    risk = result.get("risk_level", "?")
    print(f"\nSignal: {signal} | Confidence: {conf} | Risk: {risk}")


if __name__ == "__main__":
    main()