import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

import torch
from datasets import load_dataset, Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    BitsAndBytesConfig,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType
from dataclasses import dataclass
from transformers import DataCollatorWithPadding

MODEL_NAME = "Qwen/Qwen2.5-Coder-7B-Instruct"
OUTPUT_DIR = "/home/ai/qwen-coder-programming-best"
SYSTEM = "You are an expert software engineer. Write clean, correct, efficient, well-structured code with brief explanations."

EXPERT_EXAMPLES = [
    {"instruction": "Write a Python function to check if a string is a valid palindrome ignoring case and non-alphanumeric characters.", "input": "", "output": "```python\ndef is_palindrome(s: str) -> bool:\n    cleaned = ''.join(c.lower() for c in s if c.isalnum())\n    return cleaned == cleaned[::-1]\n```\n\nTime: O(n), Space: O(n). Returns True for 'A man, a plan, a canal: Panama'."},
    {"instruction": "Write a Python function using a binary search to find the index of a target in a sorted list, or -1 if not present.", "input": "", "output": "```python\ndef binary_search(arr, target):\n    low, high = 0, len(arr) - 1\n    while low <= high:\n        mid = (low + high) // 2\n        if arr[mid] == target:\n            return mid\n        elif arr[mid] < target:\n            low = mid + 1\n        else:\n            high = mid - 1\n    return -1\n```\n\nTime: O(log n), Space: O(1)."},
    {"instruction": "Write a Python class implementing a simple LRU cache with get and put operations in O(1).", "input": "", "output": "```python\nfrom collections import OrderedDict\n\nclass LRUCache:\n    def __init__(self, capacity: int):\n        self.capacity = capacity\n        self.cache = OrderedDict()\n\n    def get(self, key: int) -> int:\n        if key not in self.cache:\n            return -1\n        self.cache.move_to_end(key)\n        return self.cache[key]\n\n    def put(self, key: int, value: int) -> None:\n        if key in self.cache:\n            self.cache.move_to_end(key)\n        self.cache[key] = value\n        if len(self.cache) > self.capacity:\n            self.cache.popitem(last=False)\n```\n\nBoth operations are O(1) using OrderedDict."},
    {"instruction": "Write Python code using a queue (BFS) to find the shortest path length in a grid from (0,0) to (n-1,m-1) avoiding walls.", "input": "", "output": "```python\nfrom collections import deque\n\ndef shortest_path(grid):\n    n, m = len(grid), len(grid[0])\n    if grid[0][0] == 1 or grid[n-1][m-1] == 1:\n        return -1\n    q = deque([(0, 0, 1)])\n    grid[0][0] = 1\n    while q:\n        r, c, d = q.popleft()\n        if (r, c) == (n-1, m-1):\n            return d\n        for dr, dc in ((1,0),(-1,0),(0,1),(0,-1)):\n            nr, nc = r+dr, c+dc\n            if 0 <= nr < n and 0 <= nc < m and grid[nr][nc] == 0:\n                grid[nr][nc] = 1\n                q.append((nr, nc, d+1))\n    return -1\n```\n\nBFS on grid: O(n*m)."},
    {"instruction": "Write Python code using Dijkstra's algorithm to find the shortest path in a weighted graph.", "input": "", "output": "```python\nimport heapq\n\ndef dijkstra(graph, start):\n    dist = {node: float('inf') for node in graph}\n    dist[start] = 0\n    pq = [(0, start)]\n    while pq:\n        d, u = heapq.heappop(pq)\n        if d > dist[u]:\n            continue\n        for v, w in graph[u]:\n            nd = d + w\n            if nd < dist[v]:\n                dist[v] = nd\n                heapq.heappush(pq, (nd, v))\n    return dist\n```\n\nPriority-queue Dijkstra: O((E+V) log V)."},
    {"instruction": "Write a Python function using dynamic programming to compute the length of the longest increasing subsequence in an array.", "input": "", "output": "```python\ndef lis_length(nums):\n    import bisect\n    tails = []\n    for x in nums:\n        i = bisect.bisect_left(tails, x)\n        if i == len(tails):\n            tails.append(x)\n        else:\n            tails[i] = x\n    return len(tails)\n```\n\nPatience sorting: O(n log n) time, O(n) space."},
]

# ----------------------------------------------------------------
# Load and filter CodeAlpaca-18k
ds = load_dataset("iamtarun/python_code_instructions_18k_alpaca", split="train")

def filter_quality(ex):
    out = ex["output"] or ""
    ins = ex["instruction"] or ""
    if not ins.strip() or len(ins) > 2000:
        return False
    if len(out) < 10 or len(out) > 2000:
        return False
    if "def " not in out and "import " not in out and "class " not in out and "{" not in out and "function" not in out.lower():
        return False
    return True

ds = ds.filter(filter_quality, num_proc=4)
ds = ds.shuffle(seed=42)
N = 6000
ds = ds.select(range(min(N, len(ds))))

expert_ds = Dataset.from_list(EXPERT_EXAMPLES)

def normalize(ex):
    return {"instruction": ex["instruction"], "input_text": str(ex.get("input") or ""), "output": ex["output"]}

ds = ds.map(normalize, remove_columns=ds.column_names)
expert_ds = expert_ds.map(lambda e: {"instruction": e["instruction"], "input_text": str(e.get("input") or ""), "output": e["output"]})
ds = Dataset.from_list(expert_ds.to_list() + ds.to_list())
print(f"Final train examples: {len(ds)}")

# ----------------------------------------------------------------
def main():
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True, use_fast=True)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME, quantization_config=bnb_config, device_map="auto",
        trust_remote_code=True, torch_dtype=torch.bfloat16,
    )
    model = prepare_model_for_kbit_training(model)
    model.gradient_checkpointing_enable()
    model.enable_input_require_grads()

    lora_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=TaskType.CAUSAL_LM,
        use_rslora=True,
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    def format_chat(ex):
        user_content = ex["instruction"] + (f"\n\nInput:\n{ex['input_text']}" if ex["input_text"] else "")
        prefix = f"<|im_start|>system\n{SYSTEM}<|im_end|>\n<|im_start|>user\n{user_content}<|im_end|>\n"
        suffix = f"<|im_start|>assistant\n{ex['output']}<|im_end|>"
        return {"prefix": prefix, "suffix": suffix}

    def tokenize_with_masks(ex, max_length=2048):
        p = tokenizer(ex["prefix"], add_special_tokens=False)
        s = tokenizer(ex["suffix"], add_special_tokens=False)
        input_ids = (p["input_ids"] + s["input_ids"])[:max_length]
        labels = ([-100] * len(p["input_ids"]) + s["input_ids"])[:max_length]
        attention_mask = [1] * len(input_ids)
        return {"input_ids": input_ids, "labels": labels, "attention_mask": attention_mask}

    ds_map = ds.map(format_chat, remove_columns=ds.column_names)
    ds_map = ds_map.map(
        lambda e: tokenize_with_masks(e),
        remove_columns=ds_map.column_names,
        batched=False,
    )
    print("Sample input ids len:", len(ds_map[0]["input_ids"]))

    class CompletionOnlyDataCollator:
        def __init__(self, tokenizer):
            self.tokenizer = tokenizer
        def __call__(self, features):
            pad_id = self.tokenizer.pad_token_id
            max_len = max(len(f["input_ids"]) for f in features)
            input_ids, attention_mask, labels = [], [], []
            for f in features:
                ids, mask, lab = f["input_ids"], f["attention_mask"], f["labels"]
                pad = max_len - len(ids)
                input_ids.append(ids + [pad_id] * pad)
                attention_mask.append(mask + [0] * pad)
                labels.append(lab + [-100] * pad)
            return {
                "input_ids": torch.tensor(input_ids, dtype=torch.long),
                "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
                "labels": torch.tensor(labels, dtype=torch.long),
            }

    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=3,
        per_device_train_batch_size=8,
        gradient_accumulation_steps=4,
        warmup_steps=80,
        learning_rate=2e-4,
        lr_scheduler_type="cosine",
        weight_decay=0.0,
        fp16=False,
        bf16=True,
        max_grad_norm=1.0,
        logging_steps=20,
        save_steps=250,
        save_total_limit=2,
        optim="adamw_8bit",
        report_to="none",
        remove_unused_columns=False,
        dataloader_pin_memory=False,
        gradient_checkpointing_kwargs={"use_reentrant": False},
    )

    collator = CompletionOnlyDataCollator(tokenizer=tokenizer)

    trainer = Trainer(
        model=model,
        train_dataset=ds_map,
        args=training_args,
        data_collator=collator,
    )

    print("Starting aggressive training...")
    trainer.train()

    trainer.save_model()
    tokenizer.save_pretrained(OUTPUT_DIR)
    print(f"Model saved to {OUTPUT_DIR}")

if __name__ == "__main__":
    main()