import pandas as pd
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import streamlit as st

class CSVAnalyzer:
    @st.cache_resource
    def load_model():
        model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
        try:
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                torch_dtype=torch.float32,
                device_map="auto",
                trust_remote_code=True
            )
            return tokenizer, model
        except Exception as e:
            print(f"Initialisierungsfehler: {str(e)}")
            raise

    def __init__(self):
        self.tokenizer, self.model = CSVAnalyzer.load_model()

    def prepare_context(self, df: pd.DataFrame) -> str:
        """
        Bereitet einen allgemeinen Kontext aus dem DataFrame vor.
        """
        try:
            if df.empty:
                return "Keine Daten verfügbar"
            
            # Grundlegende Informationen
            context = "DATASET ÜBERSICHT:\n\n"
            
            # Dimensionen
            context += f"Datensatzgröße: {len(df)} Zeilen, {len(df.columns)} Spalten\n\n"
            
            # Spaltenliste mit Datentypen
            context += "SPALTENINFORMATIONEN:\n"
            for col in df.columns:
                dtype = df[col].dtype
                non_null = df[col].count()
                null_percentage = (len(df) - non_null) / len(df) * 100
                
                # Erkennung spezieller Datentypen
                if pd.api.types.is_datetime64_any_dtype(df[col]):
                    date_range = f"von {df[col].min()} bis {df[col].max()}"
                    context += f"- {col} (Datum): {date_range}, {null_percentage:.1f}% NULL\n"
                elif pd.api.types.is_numeric_dtype(df[col]):
                    context += f"- {col} (Numerisch): Min={df[col].min():.2f}, Max={df[col].max():.2f}, {null_percentage:.1f}% NULL\n"
                else:
                    unique_values = df[col].nunique()
                    context += f"- {col} (Text): {unique_values} eindeutige Werte, {null_percentage:.1f}% NULL\n"
            
            # Grundlegende Statistiken für numerische Spalten
            numeric_cols = df.select_dtypes(include=['int64', 'float64']).columns
            if not numeric_cols.empty:
                context += "\nNUMERISCHE STATISTIKEN:\n"
                stats = df[numeric_cols].describe()
                for col in numeric_cols:
                    context += f"- {col}:\n"
                    context += f"  Durchschnitt: {stats[col]['mean']:.2f}\n"
                    context += f"  Median: {stats[col]['50%']:.2f}\n"
                    context += f"  Standardabweichung: {stats[col]['std']:.2f}\n"
            
            # Zeitliche Informationen, falls vorhanden
            date_cols = df.select_dtypes(include=['datetime64']).columns
            if not date_cols.empty:
                context += "\nZEITLICHE INFORMATIONEN:\n"
                for col in date_cols:
                    context += f"- {col}:\n"
                    context += f"  Zeitspanne: von {df[col].min()} bis {df[col].max()}\n"
                    context += f"  Anzahl eindeutiger Daten: {df[col].nunique()}\n"
            
            # Stichprobe der Daten
            context += "\nDATENBEISPIELE:\n"
            sample = df.head(3).to_string()
            context += f"{sample}\n"
            
            return context
            
        except Exception as e:
            raise Exception(f"Fehler bei der Kontextvorbereitung: {str(e)}")

    def generate_response(self, context: str, query: str) -> str:
        prompt = f"""<|system|>
Du bist ein Datenanalyst, der CSV-Dateien analysiert. Deine Aufgabe ist es, Fragen über die Daten zu beantworten.
Berücksichtige dabei:
1. Die Struktur und die Arten der verfügbaren Daten
2. Statistische Informationen, falls relevant
3. Mögliche Zusammenhänge zwischen verschiedenen Spalten
4. Zeitliche Muster, falls Datumsinformationen vorhanden sind

Antworte präzise und faktenbasiert. Wenn die Frage nicht mit den verfügbaren Daten beantwortet werden kann, 
erkläre warum.

<|user|>
KONTEXT:
{context}

FRAGE:
{query}

<|assistant|>"""

        try:
            inputs = self.tokenizer(
                prompt,
                return_tensors="pt",
                truncation=True,
                max_length=2048,  # Erhöht für komplexere Analysen
                padding=True,
                return_attention_mask=True
            )
            
            with torch.no_grad():
                outputs = self.model.generate(
                    input_ids=inputs["input_ids"],
                    attention_mask=inputs["attention_mask"],
                    max_new_tokens=256,  # Erhöht für ausführlichere Antworten
                    temperature=0.7,     # Erhöht für kreativere Analysen
                    top_p=0.95,
                    repetition_penalty=1.15,
                    do_sample=True
                )
            
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            response = response.split("<|assistant|>")[-1].strip()
            
            return response
            
        except Exception as e:
            return f"Generierungsfehler: {str(e)}"

def analyze_csv(df: pd.DataFrame, query: str) -> str:
    """Hauptfunktion für die Analyse von CSV-Daten."""
    try:
        analyzer = CSVAnalyzer()
        context = analyzer.prepare_context(df)
        response = analyzer.generate_response(context, query)
        return response
    except Exception as e:
        return f"Analysefehler: {str(e)}"