# Install necessary libraries
!pip install fastapi uvicorn requests pandas scikit-learn pydantic tenacity python-json-logger --quiet

import logging
import sys
import random
import time
from collections import deque
from typing import Dict, Any, List, Optional

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
import joblib

from fastapi import FastAPI, Request, HTTPException
from pydantic import BaseModel
import uvicorn
import requests
from tenacity import retry, wait_exponential, stop_after_attempt, retry_if_exception_type
from pythonjsonlogger import jsonlogger

# Configure logging
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# Use jsonlogger for structured logging
handler = logging.StreamHandler(sys.stdout)
formatter = jsonlogger.JsonFormatter(
    '%(levelname)s %(asctime)s %(module)s %(funcName)s %(lineno)d %(message)s'
)
handler.setFormatter(formatter)
logger.addHandler(handler)

logger.info("All necessary libraries imported and logger configured.")

def create_model_state(model: Optional[Any] = None, feature_names: Optional[List[str]] = None) -> Dict[str, Any]:
    """
    Initializes a dictionary to hold the current state of the machine learning model.

    Parameters
    ----------
    model : Optional[Any], optional
        The trained machine learning model object. Defaults to None.
    feature_names : Optional[List[str]], optional
        A list of feature names the model expects. Defaults to None.

    Returns
    -------
    Dict[str, Any]
        An initialized state dictionary for the model.
    """
    state = {
        "model": model,
        "feature_names": feature_names,
        "status": "unloaded" if model is None else "loaded",
        "model_version": "1.0.0"
    }
    logger.info("Model state initialized.", extra={"state_status": state["status"]})
    return state

def load_model(state: Dict[str, Any], model_path: str = "model.joblib") -> Dict[str, Any]:
    """
    Simulates loading a machine learning model from a specified path.
    If the model doesn't exist, it trains a dummy model and saves it.

    Parameters
    ----------
    state : Dict[str, Any]
        Current state dictionary.
    model_path : str, optional
        The file path to load/save the model. Defaults to "model.joblib".

    Returns
    -------
    Dict[str, Any]
        Updated state with the loaded model and updated status.
    """
    try:
        # Try to load the model
        state["model"] = joblib.load(model_path)
        state["status"] = "loaded"
        logger.info("Model loaded successfully from path.", extra={"model_path": model_path})
    except FileNotFoundError:
        logger.warning("Model file not found. Training a dummy model.", extra={"model_path": model_path})
        # Create a dummy dataset
        X = pd.DataFrame({
            'feature_1': [random.random() * 10 for _ in range(100)],
            'feature_2': [random.random() * 5 for _ in range(100)],
            'feature_3': [random.randint(0, 100) for _ in range(100)]
        })
        y = pd.Series([1 if (x['feature_1'] + x['feature_2'] + x['feature_3']/10) > 10 else 0 for _, x in X.iterrows()])

        # Train a dummy model
        model = RandomForestClassifier(random_state=42)
        model.fit(X, y)
        state["model"] = model
        state["feature_names"] = list(X.columns)
        state["status"] = "loaded"

        # Save the trained model
        joblib.dump(model, model_path)
        logger.info("Dummy model trained and saved.", extra={"model_path": model_path, "feature_names": state["feature_names"]})
    except Exception as e:
        logger.error("Error loading model.", exc_info=True, extra={"error": str(e)})
        state["status"] = "error"

    return state

def preprocess_data(state: Dict[str, Any], raw_data: Dict[str, Any]) -> pd.DataFrame:
    """
    Simulates preprocessing input data for the model.

    Parameters
    ----------
    state : Dict[str, Any]
        Current state dictionary, containing feature names.
    raw_data : Dict[str, Any]
        Raw input data as a dictionary.

    Returns
    -------
    pd.DataFrame
        Preprocessed data as a DataFrame.
    """
    if state.get("feature_names") is None:
        logger.error("Feature names not found in state, cannot preprocess data.")
        raise ValueError("Feature names are required for preprocessing.")

    try:
        # Convert raw data to DataFrame, ensuring feature order
        input_df = pd.DataFrame([raw_data], columns=state["feature_names"])
        # In a real scenario, more complex preprocessing would happen here (e.g., scaling, encoding)
        logger.info("Data preprocessed successfully.", extra={"input_features": list(input_df.columns)})
        return input_df
    except KeyError as e:
        logger.error("Missing required feature in raw data during preprocessing.", extra={"missing_feature": str(e)})
        raise HTTPException(status_code=422, detail=f"Missing required feature: {e}")
    except Exception as e:
        logger.error("Error during data preprocessing.", exc_info=True, extra={"error": str(e)})
        raise HTTPException(status_code=500, detail="Internal server error during preprocessing.")

def predict(state: Dict[str, Any], preprocessed_data: pd.DataFrame) -> List[Any]:
    """
    Generates predictions using the loaded machine learning model.

    Parameters
    ----------
    state : Dict[str, Any]
        Current state dictionary, containing the loaded model.
    preprocessed_data : pd.DataFrame
        Preprocessed data for prediction.

    Returns
    -------
    List[Any]
        A list of predictions.
    """
    model = state.get("model")
    if model is None or state.get("status") != "loaded":
        logger.error("Model not loaded for prediction.", extra={"model_status": state.get("status")})
        raise HTTPException(status_code=503, detail="Model is not loaded or ready for predictions.")

    try:
        predictions = model.predict(preprocessed_data)
        logger.info("Prediction successful.", extra={"num_predictions": len(predictions)})
        return predictions.tolist()
    except Exception as e:
        logger.error("Error during model prediction.", exc_info=True, extra={"error": str(e)})
        raise HTTPException(status_code=500, detail="Internal server error during prediction.")

def create_fastapi_app(model_state: Dict[str, Any]) -> FastAPI:
    """
    Initializes a FastAPI application and loads the model into its state.

    Parameters
    ----------
    model_state : Dict[str, Any]
        The initial model state dictionary.

    Returns
    -------
    FastAPI
        An initialized FastAPI application instance.
    """
    app = FastAPI(
        title="ML Model Serving API",
        description="API for serving a machine learning model",
        version=model_state.get("model_version", "1.0.0")
    )

    # Store the model state in the app's state
    app.state.model_state = model_state

    @app.get("/health")
    async def health_check():
        logger.info("Health check requested.")
        return {"status": "ok", "model_status": app.state.model_state["status"]}

    logger.info("FastAPI application initialized.")
    return app

class PredictionRequest(BaseModel):
    feature_1: float
    feature_2: float
    feature_3: int

    class Config:
        schema_extra = {
            "example": {
                "feature_1": 5.5,
                "feature_2": 2.1,
                "feature_3": 75
            }
        }

logger.info("Pydantic model 'PredictionRequest' defined.")

def add_prediction_endpoint(app: FastAPI) -> FastAPI:
    """
    Adds a /predict endpoint to the FastAPI application.

    Parameters
    ----------
    app : FastAPI
        The FastAPI application instance.

    Returns
    -------
    FastAPI
        The FastAPI application instance with the added prediction endpoint.
    """

    @app.post("/predict")
    async def predict_endpoint(request_data: PredictionRequest, request: Request):
        logger.info("Prediction request received.", extra={"request_data": request_data.dict()})
        model_state = request.app.state.model_state

        try:
            # Preprocess data
            preprocessed_data = preprocess_data(model_state, request_data.dict())
            # Make prediction
            prediction = predict(model_state, preprocessed_data)
            logger.info("Prediction successful, sending response.", extra={"prediction_result": prediction})
            return {"prediction": prediction}
        except HTTPException as e:
            logger.error("HTTPException during prediction endpoint.", extra={"status_code": e.status_code, "detail": e.detail})
            raise e
        except Exception as e:
            logger.error("Unexpected error during prediction endpoint.", exc_info=True, extra={"error": str(e)})
            raise HTTPException(status_code=500, detail="Internal server error.")

    logger.info("Prediction endpoint '/predict' added to FastAPI application.")
    return app

def run_fastapi_app(app: FastAPI, host: str = '0.0.0.0', port: int = 8000) -> None:
    """
    Runs the FastAPI application using Uvicorn.
    Note: Running a web server directly in a Colab cell might block it.
    For interactive testing, consider running in a separate thread/process or exposing via ngrok.

    Parameters
    ----------
    app : FastAPI
        The FastAPI application instance.
    host : str, optional
        The host address to bind to. Defaults to '0.0.0.0'.
    port : int, optional
        The port number to listen on. Defaults to 8000.
    """
    logger.info("Attempting to run FastAPI application with Uvicorn.", extra={"host": host, "port": port})
    try:
        # This will block the cell. For non-blocking, consider threading/multiprocessing or ngrok.
        uvicorn.run(app, host=host, port=port)
        logger.info("FastAPI application stopped.")
    except Exception as e:
        logger.error("Failed to run FastAPI application with Uvicorn.", exc_info=True, extra={"error": str(e)})

from tenacity import wait_random

@retry(
    wait=wait_exponential(multiplier=1, min=2, max=10) + wait_random(0, 1),
    stop=stop_after_attempt(5),
    retry=retry_if_exception_type(requests.exceptions.RequestException)
)
def send_prediction_request(url: str, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """
    Sends an HTTP POST request to the /predict endpoint with retry logic and jitter.

    Parameters
    ----------
    url : str
        The base URL of the FastAPI application.
    payload : Dict[str, Any]
        The data to be sent in the request body.

    Returns
    -------
    Optional[Dict[str, Any]]
        The JSON response from the API, or None if failed.
    """
    try:
        prediction_url = f"{url}/predict"
        logger.debug("Sending prediction request.", extra={"url": prediction_url, "payload": payload})
        response = requests.post(prediction_url, json=payload, timeout=5) # Add a timeout
        response.raise_for_status()  # Raise an exception for HTTP errors (4xx or 5xx)
        logger.info("Prediction request successful.", extra={"status_code": response.status_code})
        return response.json()
    except requests.exceptions.HTTPError as e:
        logger.error("HTTP error during prediction request.", extra={"status_code": e.response.status_code, "response_text": e.response.text})
        raise  # Re-raise to trigger tenacity retry if it's a transient error
    except requests.exceptions.ConnectionError as e:
        logger.warning("Connection error during prediction request, retrying.", exc_info=True, extra={"error": str(e)})
        raise  # Re-raise to trigger tenacity retry
    except requests.exceptions.Timeout as e:
        logger.warning("Timeout during prediction request, retrying.", exc_info=True, extra={"error": str(e)})
        raise  # Re-raise to trigger tenacity retry
    except requests.exceptions.RequestException as e:
        logger.error("General request error during prediction request.", exc_info=True, extra={"error": str(e)})
        raise  # Re-raise to trigger tenacity retry
    except Exception as e:
        logger.error("Unexpected error during prediction request.", exc_info=True, extra={"error": str(e)})
        return None # For non-retryable errors, return None

# 1. Initialize model state and load a dummy model
initial_model_state = create_model_state()
model_state = load_model(initial_model_state)

# 2. Create and configure the FastAPI application
app_instance = create_fastapi_app(model_state)
app_instance = add_prediction_endpoint(app_instance)

logger.info("FastAPI app and model are prepared for conceptual serving.")

# In a real Colab scenario, you would often use ngrok to expose your local server:
# !pip install pyngrok --quiet
# from pyngrok import ngrok
# import nest_asyncio
# nest_asyncio.apply()

# # Start ngrok and get the public URL
# ngrok_tunnel = ngrok.connect(8000)
# public_url = ngrok_tunnel.public_url
# logger.info(f"FastAPI app exposed publicly at: {public_url}")

# # Run Uvicorn in a separate thread/process for non-blocking execution
# import threading
# thread = threading.Thread(target=run_fastapi_app, args=(app_instance, '0.0.0.0', 8000))
# thread.start()

# For this demonstration, we'll assume the app is running and interact with a local mock URL.
BASE_URL = "http://127.0.0.1:8000" # If using ngrok, this would be `public_url`

print(f"--- Simulating interactions with the FastAPI server at {BASE_URL} ---")

# Test health check
print("\n--- Testing Health Check ---")
try:
    health_response = requests.get(f"{BASE_URL}/health")
    health_response.raise_for_status()
    print(f"Health Check Success: {health_response.json()}")
except requests.exceptions.RequestException as e:
    print(f"Health Check Failed (Expected if server not truly running): {e}")
    logger.warning("Health check failed, likely because Uvicorn is not actively running in this Colab environment.")

# Generate some sample data for prediction
sample_data = [
    {"feature_1": 6.2, "feature_2": 3.0, "feature_3": 80},
    {"feature_1": 3.1, "feature_2": 1.5, "feature_3": 25},
    {"feature_1": 8.9, "feature_2": 4.2, "feature_3": 95},
    {"feature_1": 1.5, "feature_2": 0.7, "feature_3": 10},
    {"feature_1": 7.0, "feature_2": 3.5, "feature_3": 60}
]

predictions_results = []

print("\n--- Sending Prediction Requests ---")
for i, payload in enumerate(sample_data):
    print(f"Sending request {i+1} with payload: {payload}")
    try:
        # Call the send_prediction_request function directly for demonstration
        # In a real scenario, this would be an actual HTTP call to a running server.
        # For this Colab simulation, we'll directly call the endpoint logic for prediction.

        # NOTE: This bypasses the actual HTTP request but demonstrates the logic.
        # If you've got ngrok running and the Uvicorn thread active, use send_prediction_request.

        # For demonstration without a live server, we call the endpoint logic directly:
        # Simulate FastAPI's handling of the request
        pydantic_request = PredictionRequest(**payload)

        # Create a mock Request object for accessing app.state
        class MockAppState:
            def __init__(self, model_st):
                self.model_state = model_st
        class MockRequest:
            def __init__(self, app_state):
                self.app = MockAppState(app_state)

        mock_request = MockRequest(app_instance.state.model_state)

        # Call the internal endpoint logic
        async def internal_predict_call():
            return await app_instance.post("/predict")(pydantic_request, mock_request) # Call the decorated function

        # Since we are not in an async context, we simulate the await by directly getting the result.
        # This requires more setup, so for simplicity, let's make it more direct for now:
        # We will directly call preprocess_data and predict, rather than simulating the full FastAPI route.

        preprocessed = preprocess_data(app_instance.state.model_state, payload)
        prediction = predict(app_instance.state.model_state, preprocessed)

        result = {"prediction": prediction}
        predictions_results.append({"payload": payload, "prediction": result["prediction"]})
        print(f"  -> Prediction: {result}")

        # If you were actually sending HTTP requests to a live server:
        # response_json = send_prediction_request(BASE_URL, payload)
        # if response_json:
        #     predictions_results.append({"payload": payload, "prediction": response_json["prediction"]})
        #     print(f"  -> Prediction: {response_json}")
        # else:
        #     print("  -> Prediction failed after retries.")

    except Exception as e:
        print(f"  -> Error during prediction for payload {payload}: {e}")
        logger.error("Error during prediction simulation.", exc_info=True, extra={"payload": payload, "error": str(e)})


# Display results in a DataFrame
print("\n--- Summary of Predictions ---")
predictions_df = pd.DataFrame(predictions_results)
display(predictions_df)

logger.info("Demonstration of prediction requests completed.")

import matplotlib.pyplot as plt
import seaborn as sns

# Flatten the list of lists for predictions_results
all_predictions = [pred for item in predictions_results for pred in item['prediction']]

if all_predictions:
    predictions_series = pd.Series(all_predictions, name="prediction_class")

    plt.figure(figsize=(8, 6))
    sns.countplot(x=predictions_series, palette='viridis', hue=predictions_series, legend=False)
    plt.title('Distribution of Predicted Classes')
    plt.xlabel('Predicted Class')
    plt.ylabel('Count')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.xticks(ticks=[0, 1], labels=['Class 0', 'Class 1']) # Assuming binary classification
    plt.tight_layout()
    plt.show()
    logger.info("Prediction distribution visualized.")
else:
    print("No predictions to visualize.")
    logger.warning("No predictions available for visualization.")