#!/usr/bin/env python3
"""
XYFLOW Dual-Copy Thermal Bypass for HuggingFaceTB/SmolLM2-135M-Instruct
Fixed: proper passing of input_ids for each chunk.
"""

import os
import time
import threading
import psutil
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# =============================================================================
# CONFIGURATION
# =============================================================================
MODEL_NAME = "HuggingFaceTB/SmolLM2-135M-Instruct"
PROMPT = "Explain the theory of relativity in simple terms."
MAX_NEW_TOKENS = 120
SWITCH_INTERVAL = 6          # tokens per chunk
THERMAL_THRESHOLD = 80.0     # °C – early switch if exceeded

# Physical core IDs (avoid hyperthreads)
CORE_A = 0
CORE_B = 2

# =============================================================================
# SYSTEM UTILITIES
# =============================================================================

def set_affinity(pid, core_id):
    try:
        os.sched_setaffinity(pid, {core_id})
        return True
    except (AttributeError, OSError):
        try:
            psutil.Process(pid).cpu_affinity([core_id])
            return True
        except AttributeError:
            print("Warning: CPU affinity not supported.")
            return False

def get_cpu_temperature():
    try:
        temps = psutil.sensors_temperatures()
        for sensor in temps.values():
            if sensor:
                return sensor[0].current
    except:
        pass
    return 0.0

# =============================================================================
# THERMAL MONITOR
# =============================================================================

class ThermalMonitor:
    def __init__(self, interval=0.5):
        self.interval = interval
        self.running = False
        self.temps = {"core0": [], "core1": [], "time": []}
        self.thread = None

    def start(self):
        self.running = True
        self.thread = threading.Thread(target=self._run)
        self.thread.daemon = True
        self.thread.start()

    def stop(self):
        self.running = False
        if self.thread:
            self.thread.join()

    def _run(self):
        start = time.time()
        while self.running:
            temp = get_cpu_temperature()
            self.temps["core0"].append(temp)
            self.temps["core1"].append(temp)   # package temp for both
            self.temps["time"].append(time.time() - start)
            time.sleep(self.interval)

# =============================================================================
# DUAL-COPY ENGINE
# =============================================================================

class DualCopyInference:
    def __init__(self):
        # Force single‑threaded inference for precise core affinity
        os.environ["OMP_NUM_THREADS"] = "1"
        os.environ["MKL_NUM_THREADS"] = "1"
        torch.set_num_threads(1)

        print("Loading Tokenizer...")
        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        print(f"Loading Model Copy A (pinned to Core {CORE_A})...")
        self.model_a = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME, torch_dtype=torch.float32, low_cpu_mem_usage=True
        )
        self.model_a.to('cpu')

        print(f"Loading Model Copy B (pinned to Core {CORE_B})...")
        self.model_b = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME, torch_dtype=torch.float32, low_cpu_mem_usage=True
        )
        self.model_b.to('cpu')

        self.active_model = self.model_a
        self.active_core = CORE_A
        self.past_kv = None
        self.attention_mask = None

    def _switch_model(self):
        if self.active_model is self.model_a:
            self.active_model = self.model_b
            self.active_core = CORE_B
        else:
            self.active_model = self.model_a
            self.active_core = CORE_A
        set_affinity(0, self.active_core)
        print(f"\n[SWITCH] -> Model {'B' if self.active_core == CORE_B else 'A'} on Core {self.active_core}")

    def _generate_chunk(self, start_ids, past_kv, attention_mask, max_new_tokens):
        """
        Generate up to max_new_tokens starting from start_ids.
        start_ids: tensor of shape (1, seq_len) – either full prompt or last token.
        past_kv, attention_mask: cached state from previous chunks.
        Returns: (list_of_token_ids, updated_past_kv, updated_attention_mask)
        """
        generated_tokens = []
        current_ids = start_ids
        current_past = past_kv
        current_mask = attention_mask

        for _ in range(max_new_tokens):
            with torch.no_grad():
                outputs = self.active_model(
                    input_ids=current_ids,
                    past_key_values=current_past,
                    attention_mask=current_mask,
                    use_cache=True
                )
            logits = outputs.logits[:, -1, :]
            next_token = torch.argmax(logits, dim=-1, keepdim=True)
            generated_tokens.append(next_token.item())

            # update cache and mask
            current_past = outputs.past_key_values
            current_ids = next_token
            if current_mask is not None:
                current_mask = torch.cat([current_mask, torch.ones((1, 1), dtype=torch.long)], dim=1)
            else:
                current_mask = torch.ones((1, 1), dtype=torch.long)

        return generated_tokens, current_past, current_mask

    def generate(self, prompt, max_tokens=MAX_NEW_TOKENS, switch_interval=SWITCH_INTERVAL):
        print(f"\nStarting Generation (switch every {switch_interval} tokens)")
        print("-" * 60)

        # Initial tokenization
        inputs = self.tokenizer(prompt, return_tensors="pt")
        input_ids = inputs["input_ids"]
        attention_mask = inputs.get("attention_mask", None)
        past_kv = None

        all_tokens = []
        remaining = max_tokens
        token_count_since_switch = 0
        last_token_id = None

        # Start thermal monitor
        monitor = ThermalMonitor()
        monitor.start()

        # Pin to first core and set active model
        set_affinity(0, CORE_A)
        self.active_model = self.model_a
        self.active_core = CORE_A

        try:
            while remaining > 0:
                chunk_size = min(switch_interval, remaining)

                # Thermal emergency switch
                temp = get_cpu_temperature()
                if temp >= THERMAL_THRESHOLD and self.active_core is not None:
                    print(f"\n[THERMAL TRIGGER] {temp:.1f}°C >= {THERMAL_THRESHOLD}°C. Forcing early switch.")
                    self._switch_model()
                    token_count_since_switch = 0  # reset to avoid immediate switch back

                # Determine starting token(s) for this chunk
                if past_kv is None:
                    # First chunk: use the entire prompt
                    start_ids = input_ids
                else:
                    # Subsequent chunks: use only the last generated token
                    start_ids = torch.tensor([[last_token_id]], dtype=torch.long)

                print(f"Generating {chunk_size} tokens on Core {self.active_core}...", end="", flush=True)
                start_time = time.time()

                chunk_tokens, past_kv, attention_mask = self._generate_chunk(
                    start_ids, past_kv, attention_mask, chunk_size
                )

                all_tokens.extend(chunk_tokens)
                last_token_id = chunk_tokens[-1]      # remember for next chunk
                remaining -= chunk_size
                token_count_since_switch += chunk_size

                elapsed = time.time() - start_time
                print(f" done in {elapsed:.2f}s ({chunk_size/elapsed:.2f} tok/s)")

                # Switch if we've completed the interval and more tokens remain
                if remaining > 0 and token_count_since_switch >= switch_interval:
                    self._switch_model()
                    token_count_since_switch = 0

        finally:
            monitor.stop()

        # Decode final output
        full_ids = torch.cat([inputs["input_ids"], torch.tensor([all_tokens])], dim=1)
        output_text = self.tokenizer.decode(full_ids[0], skip_special_tokens=True)

        print("-" * 60)
        print(f"Generation Complete. Total Tokens: {len(all_tokens)}")
        print(f"Peak Temperature: {max(monitor.temps['core0']):.1f}°C")
        return output_text, monitor.temps

# =============================================================================
# BASELINE
# =============================================================================

def baseline_single_copy(prompt, max_tokens=MAX_NEW_TOKENS):
    print("\n" + "="*60)
    print("RUNNING BASELINE (Single Copy, No Switching)")
    print("="*60)

    # Use multiple threads for normal inference
    os.environ["OMP_NUM_THREADS"] = "4"
    torch.set_num_threads(4)

    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    inputs = tokenizer(prompt, return_tensors="pt")
    monitor = ThermalMonitor()
    monitor.start()
    start_time = time.time()

    outputs = model.generate(
        **inputs,
        max_new_tokens=max_tokens,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id,
        use_cache=True
    )

    elapsed = time.time() - start_time
    monitor.stop()
    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"Baseline complete in {elapsed:.2f}s")
    return text, monitor.temps

# =============================================================================
# MAIN
# =============================================================================

if __name__ == "__main__":
    print("="*60)
    print("XYFLOW DUAL-COPY THERMAL EXPERIMENT")
    print(f"Model: {MODEL_NAME}")
    print(f"Switch Interval: {SWITCH_INTERVAL} tokens")
    print(f"Thermal Threshold: {THERMAL_THRESHOLD}°C")
    print("="*60)

    # 1. Baseline
    baseline_text, baseline_temps = baseline_single_copy(PROMPT)
    print(f"\nBaseline Output:\n{baseline_text[:200]}...")
    print(f"Baseline Peak Temp: {max(baseline_temps['core0']):.1f}°C")

    # 2. Dual‑Copy
    print("\n" + "="*60)
    print("INITIALIZING DUAL-COPY ENGINE...")
    print("="*60)

    engine = DualCopyInference()
    dual_text, dual_temps = engine.generate(PROMPT)
    print(f"\nDual-Copy Output:\n{dual_text[:200]}...")
    print(f"Dual-Copy Peak Temp: {max(dual_temps['core0']):.1f}°C")

    # 3. Conclusion
    print("\n" + "="*60)
    print("EXPERIMENT CONCLUSION")
    print("="*60)
    baseline_peak = max(baseline_temps['core0'])
    dual_peak = max(dual_temps['core0'])
    if dual_peak < baseline_peak:
        print(f"✅ SUCCESS: Peak temperature dropped from {baseline_peak:.1f}°C to {dual_peak:.1f}°C.")
        print("   The CPU has more thermal headroom, sustaining higher boost clocks.")
    else:
        print(f"⚠️  NEUTRAL: Similar peak temps ({baseline_peak:.1f}°C vs {dual_peak:.1f}°C).")
        print("   Try adjusting SWITCH_INTERVAL or using different physical cores.")