from dataclasses import dataclass
from typing import Dict, List, Optional
import time
import json
import subprocess
import statistics
from pathlib import Path
from .utils import get_memory_usage, MemorySnapshot
import psutil

@dataclass
class QuantizationMetrics:
    # Base Model Info
    original_size: float  # Size in MB
    quantized_size: float  # Size in MB
    compression_ratio: float  # original/quantized

    # Performance Metrics
    inference_speed: float  # tokens/second
    memory_usage: float  # MB
    gpu_memory_usage: Optional[float]  # MB

    # Accuracy Metrics
    perplexity: float
    accuracy_score: float  # Based on benchmark dataset
    relative_quality: float  # Compared to full precision (0-1)

    # Resource Impact
    cpu_utilization: float
    power_consumption: Optional[float]
    temperature_impact: Optional[float]

    # Stability Metrics
    numerical_stability: float  # 0-1 score
    error_rate: float
    nan_infinity_occurrences: int

@dataclass
class BenchmarkExample:
    """Example for benchmarking model performance.

    Attributes:
        prompt: Input text to the model
        reference: Expected or reference output
        category: Type of task (e.g., "general", "technical", "creative")
    """
    prompt: str
    reference: str
    category: str

class QuantizationBenchmark:
    def __init__(
        self,
        model_name: str,
        reference_model: Optional[str] = None,
        benchmark_dataset: Optional[List[BenchmarkExample]] = None
    ):
        self.model_name = model_name
        self.reference_model = reference_model
        self.benchmark_dataset = benchmark_dataset or self._get_default_benchmark_dataset()

    def _get_default_benchmark_dataset(self) -> List[BenchmarkExample]:
        """Provide a default set of benchmark examples."""
        return [
            BenchmarkExample(
                prompt="Explain how a car engine works.",
                reference="A car engine is an internal combustion engine that converts fuel into mechanical energy...",
                category="technical"
            ),
            BenchmarkExample(
                prompt="Write a poem about spring.",
                reference="Gentle breeze and budding flowers...",
                category="creative"
            ),
            BenchmarkExample(
                prompt="What is the capital of France?",
                reference="The capital of France is Paris.",
                category="factual"
            ),
            # Add more examples covering different aspects
        ]

    def _get_model_sizes(self) -> tuple[float, float]:
        """Get original and quantized model sizes."""
        try:
            cmd = ["cortex", "models", "get", self.model_name, "--json"]
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            model_info = json.loads(result.stdout)

            # Get original size from model info (if available)
            original_size = model_info.get("original_size", 0)
            quantized_size = model_info.get("size", 0)

            return original_size / (1024 * 1024), quantized_size / (1024 * 1024)  # Convert to MB
        except Exception as e:
            print(f"Error getting model sizes: {e}")
            return 0.0, 0.0

    def _measure_inference_speed(self) -> float:
        """Measure inference speed with standard prompt."""
        standard_prompt = "Tell me a short story about a robot."

        try:
            start_time = time.time()
            result = subprocess.run(
                ["cortex", "run", "-d", self.model_name, "--json"],
                input=json.dumps({"prompt": standard_prompt}),
                capture_output=True,
                text=True,
                check=True
            )
            end_time = time.time()

            response = json.loads(result.stdout)
            tokens_generated = response.get("usage", {}).get("completion_tokens", 0)

            return tokens_generated / (end_time - start_time)
        except Exception as e:
            print(f"Error measuring inference speed: {e}")
            return 0.0

    def _calculate_perplexity(self, text: str) -> float:
        """Calculate perplexity for given text."""
        try:
            cmd = ["cortex", "run", "-d", self.model_name, "--perplexity", "--json"]
            result = subprocess.run(
                cmd,
                input=json.dumps({"prompt": text}),
                capture_output=True,
                text=True,
                check=True
            )
            response = json.loads(result.stdout)
            return float(response.get("perplexity", float('inf')))
        except Exception:
            return float('inf')

    def _evaluate_accuracy(self) -> tuple[float, float]:
        """Evaluate model accuracy using benchmark dataset."""
        accuracy_scores = []
        stability_issues = 0

        for example in self.benchmark_dataset:
            try:
                result = subprocess.run(
                    ["cortex", "run", "-d", self.model_name, "--json"],
                    input=json.dumps({"prompt": example.prompt}),
                    capture_output=True,
                    text=True,
                    check=True
                )

                response = json.loads(result.stdout)
                generated_text = response.get("choices", [{}])[0].get("text", "")

                # Calculate similarity score (simplified)
                similarity = self._calculate_similarity(generated_text, example.reference)
                accuracy_scores.append(similarity)

                # Check for numerical instability
                if "nan" in generated_text.lower() or "infinity" in generated_text.lower():
                    stability_issues += 1

            except Exception as e:
                print(f"Error during accuracy evaluation: {e}")
                stability_issues += 1

        return (
            statistics.mean(accuracy_scores) if accuracy_scores else 0.0,
            stability_issues
        )

    def _calculate_similarity(self, text1: str, text2: str) -> float:
        """Calculate similarity between two texts (simplified)."""
        # This is a very basic implementation
        # In practice, you might want to use more sophisticated metrics
        words1 = set(text1.lower().split())
        words2 = set(text2.lower().split())

        if not words1 or not words2:
            return 0.0

        intersection = words1.intersection(words2)
        union = words1.union(words2)

        return len(intersection) / len(union)

    def _measure_resource_usage(self) -> tuple[float, Optional[float], Optional[float]]:
        """Measure CPU, power, and temperature impact."""
        start_usage = get_memory_usage()
        cpu_readings = []
        temp_readings = []

        # Monitor for a short period
        for _ in range(5):
            cpu_percent = psutil.cpu_percent(interval=1)
            cpu_readings.append(cpu_percent)

            try:
                temperatures = psutil.sensors_temperatures()
                if temperatures and 'coretemp' in temperatures:
                    temp_readings.append(temperatures['coretemp'][0].current)
            except Exception:
                pass

        # Get GPU power consumption if available
        try:
            gpu_stats = gpustat.GPUStatCollection.new_query()
            power_consumption = gpu_stats[0].power_draw if gpu_stats else None
        except Exception:
            power_consumption = None

        return (
            statistics.mean(cpu_readings),
            power_consumption,
            statistics.mean(temp_readings) if temp_readings else None
        )

    def benchmark_quantization(self) -> QuantizationMetrics:
        """Run comprehensive quantization benchmarks."""
        # Get model sizes
        original_size, quantized_size = self._get_model_sizes()
        compression_ratio = original_size / quantized_size if quantized_size > 0 else 0

        # Measure inference performance
        inference_speed = self._measure_inference_speed()

        # Get memory usage
        memory_snapshot = get_memory_usage()

        # Calculate perplexity
        perplexity = self._calculate_perplexity("This is a test sentence.")

        # Evaluate accuracy and stability
        accuracy_score, stability_issues = self._evaluate_accuracy()

        # Calculate relative quality (if reference model provided)
        relative_quality = 1.0  # Default to 1.0 if no reference
        if self.reference_model:
            try:
                reference_accuracy, _ = self._evaluate_accuracy()
                relative_quality = accuracy_score / reference_accuracy
            except Exception:
                pass

        # Measure resource usage
        cpu_util, power_consumption, temp_impact = self._measure_resource_usage()

        return QuantizationMetrics(
            original_size=original_size,
            quantized_size=quantized_size,
            compression_ratio=compression_ratio,

            inference_speed=inference_speed,
            memory_usage=memory_snapshot.ram_used,
            gpu_memory_usage=memory_snapshot.gpu_memory_used,

            perplexity=perplexity,
            accuracy_score=accuracy_score,
            relative_quality=relative_quality,

            cpu_utilization=cpu_util,
            power_consumption=power_consumption,
            temperature_impact=temp_impact,

            numerical_stability=1.0 - (stability_issues / len(self.benchmark_dataset)),
            error_rate=stability_issues / len(self.benchmark_dataset),
            nan_infinity_occurrences=stability_issues
        )
