#!/usr/bin/env python3
"""
XYFLOW Dual-Copy Thermal Bypass for HuggingFaceTB/SmolLM2-135M-Instruct
Theory: Two physical copies of the weights act as two thermal variables (A, B).
Switching between them creates a 2-cycle limit cycle in the temperature phase space,
preventing the CPU from hitting the throttling separatrix.
"""

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

# =============================================================================
# CONFIGURATION (Tune these for your hardware)
# =============================================================================
MODEL_NAME = "HuggingFaceTB/SmolLM2-135M-Instruct"
PROMPT = "Explain the theory of relativity in simple terms."
MAX_NEW_TOKENS = 120
SWITCH_INTERVAL = 6          # Switch models every N generated tokens
THERMAL_THRESHOLD = 80.0     # °C - If active core exceeds this, force switch early

# CPU Core Pinning (Linux/physical cores). Adjust to your CPU topology.
# Example: Core 0 and Core 2 are physical (skip hyperthreads 1 and 3)
CORE_A = 0
CORE_B = 2

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

def set_affinity(pid, core_id):
    """Pin a process/thread to a specific CPU core (Linux)."""
    try:
        os.sched_setaffinity(pid, {core_id})
        return True
    except (AttributeError, OSError):
        # Fallback for Windows / unsupported
        try:
            psutil.Process(pid).cpu_affinity([core_id])
            return True
        except AttributeError:
            print("Warning: CPU affinity not supported on this OS.")
            return False

def get_cpu_temperature():
    """Retrieve current CPU package temperature."""
    try:
        temps = psutil.sensors_temperatures()
        if 'coretemp' in temps:
            # Intel/AMD coretemp
            return temps['coretemp'][0].current
        elif 'cpu_thermal' in temps:
            # Some ARM/Raspberry Pi
            return temps['cpu_thermal'][0].current
        else:
            # Fallback: try first available sensor
            for sensor in temps.values():
                if sensor:
                    return sensor[0].current
    except (AttributeError, KeyError, IndexError):
        pass
    return 0.0  # Cannot read temp

# =============================================================================
# THERMAL MONITOR (Background thread)
# =============================================================================

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:
            # Note: psutil doesn't easily give per-core temp on most hardware.
            # We read package temp and treat it as the shared thermal reservoir.
            # For true per-core, you'd need RAPL or vendor-specific tools.
            temp = get_cpu_temperature()
            self.temps["core0"].append(temp)
            self.temps["core1"].append(temp)  # We treat both as package temp for demo
            self.temps["time"].append(time.time() - start)
            time.sleep(self.interval)

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

class DualCopyInference:
    def __init__(self):
        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("Loading Model Copy A (pinned to Core {})...".format(CORE_A))
        self.model_a = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME, torch_dtype=torch.float32, low_cpu_mem_usage=True
        )
        self.model_a.to('cpu')
        # Force single-threaded inference so affinity pinning actually works
        torch.set_num_threads(1)

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

        # Pin the initial process (though generation will spawn threads, 
        # setting OMP_NUM_THREADS=1 ensures it stays on the pinned core)
        os.environ["OMP_NUM_THREADS"] = "1"
        os.environ["MKL_NUM_THREADS"] = "1"

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

    def _switch_model(self):
        """Swap active model and core."""
        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
        
        # Re-pin the current thread (and consequently the torch ops) 
        # to the new physical core.
        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, input_ids, past_kv, attention_mask, max_new_tokens):
        """Manual autoregressive step for precise cache control across models."""
        generated_tokens = []
        current_ids = input_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 for the next iteration
            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 with Switch Interval: {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

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

        try:
            # Initial affinity for Model A
            set_affinity(0, CORE_A)
            self.active_model = self.model_a
            self.active_core = CORE_A

            while remaining > 0:
                # Determine chunk size (switch interval, but not exceeding remaining)
                chunk_size = min(switch_interval, remaining)
                
                # Check thermal pressure: If active core is too hot, force switch early
                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()
                    # Reset counter so we don't immediately switch back
                    token_count_since_switch = 0

                # Generate on current active model
                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(
                    input_ids if past_kv is None else None, 
                    past_kv, 
                    attention_mask,
                    chunk_size
                )
                
                # Update input_ids for subsequent chunks
                if past_kv is not None:
                    input_ids = None  # We only feed the new token ids from now on

                all_tokens.extend(chunk_tokens)
                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 hit the interval, and there is more to generate
                if remaining > 0 and token_count_since_switch >= switch_interval:
                    self._switch_model()
                    token_count_since_switch = 0

        finally:
            monitor.stop()

        # Decode result
        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"Temperature Log: {monitor.temps['core0'][-1]:.1f}°C (peak)")
        return output_text, monitor.temps

# =============================================================================
# BASELINE (Single Copy) FOR COMPARISON
# =============================================================================

def baseline_single_copy(prompt, max_tokens=MAX_NEW_TOKENS):
    """Run standard inference on a single model copy (no switching)."""
    print("\n" + "="*60)
    print("RUNNING BASELINE (Single Copy, No Switching)")
    print("="*60)
    
    # Reset thread count
    os.environ["OMP_NUM_THREADS"] = "4"  # Let it use all cores for speed (normal mode)
    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 EXECUTION
# =============================================================================

if __name__ == "__main__":
    # IMPORTANT: Run this script with `python script.py`. 
    # On Linux, you may need sudo for sched_setaffinity, but usually not.

    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. Run Baseline (standard inference)
    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. Run Dual-Copy Seesaw
    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. Analysis / XYFLOW 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: The dual-copy seesaw lowered the peak temperature from {baseline_peak:.1f}°C to {dual_peak:.1f}°C.")
        print("   The CPU has higher thermal headroom, allowing sustained boost clocks.")
        print("   This validates the XYFLOW 2-Cycle Limit Cycle theory.")
    else:
        print(f"⚠️  NEUTRAL: Peak temp is similar ({baseline_peak:.1f}°C vs {dual_peak:.1f}°C).")
        print("   Try reducing SWITCH_INTERVAL or adjusting CORE_A/CORE_B to physical cores.")