import torch
import gradio as gr
import pandas as pd
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
import json
from datetime import datetime

class LegalAISystem:
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.tokenizer = AutoTokenizer.from_pretrained('nlpaueb/legal-bert-base-uncased')
        self.model = AutoModelForSequenceClassification.from_pretrained('nlpaueb/legal-bert-base-uncased')
        self.model.to(self.device)
        self.label_encoder = LabelEncoder()
        self.case_history = []

    def preprocess_data(self, text):
        # Clean and normalize text
        text = str(text).lower().strip()
        # Add more preprocessing steps as needed
        return text

    def extract_features(self, text):
        # Tokenize and prepare features
        inputs = self.tokenizer(
            text,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt"
        ).to(self.device)
        return inputs

    def predict_outcome(self, case_text):
        # Preprocess input
        processed_text = self.preprocess_data(case_text)

        # Extract features
        features = self.extract_features(processed_text)

        # Make prediction
        with torch.no_grad():
            outputs = self.model(**features)
            predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)

        # Get prediction probabilities
        probs = predictions.cpu().numpy()[0]

        # Store in case history
        self.case_history.append({
            'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            'case_text': case_text,
            'prediction_probs': probs.tolist()
        })

        return {
            'Favorable': float(probs[1]),
            'Unfavorable': float(probs[0])
        }

    def analyze_precedents(self, case_text):
        # Implement similarity search for relevant precedents
        # This is a simplified version
        return ["Precedent 1: Smith v. Jones (2019)",
                "Precedent 2: Brown v. State (2020)"]

    def generate_report(self, case_text, prediction, precedents):
        report = f"""
        Legal Case Analysis Report
        ========================
        Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

        Case Summary:
        {case_text[:500]}...

        Prediction:
        - Favorable Outcome: {prediction['Favorable']:.2%}
        - Unfavorable Outcome: {prediction['Unfavorable']:.2%}

        Relevant Precedents:
        {chr(10).join(precedents)}

        Note: This is an AI-generated analysis and should be reviewed by legal professionals.
        """
        return report

def create_gradio_interface():
    legal_ai = LegalAISystem()

    def process_case(case_text):
        # Analyze case
        prediction = legal_ai.predict_outcome(case_text)
        precedents = legal_ai.analyze_precedents(case_text)
        report = legal_ai.generate_report(case_text, prediction, precedents)

        # Create visualization data
        prob_chart = {
            "Favorable": prediction['Favorable'],
            "Unfavorable": prediction['Unfavorable']
        }

        return (
            report,
            prob_chart,
            f"Confidence: {max(prediction.values()):.2%}"
        )

    # Create Gradio interface
    iface = gr.Interface(
        fn=process_case,
        inputs=[
            gr.Textbox(label="Enter Case Details", lines=10)
        ],
        outputs=[
            gr.Textbox(label="Analysis Report", lines=10),
            gr.Label(label="Outcome Probabilities"),
            gr.Textbox(label="Model Confidence")
        ],
        title="AI Legal Case Analysis System",
        description="Enter case details to get real-time analysis and predictions."
    )

    return iface

# Launch the interface
if __name__ == "__main__":
    interface = create_gradio_interface()
    interface.launch(share=True, debug=True)