# Install necessary libraries
%pip install pandas numpy matplotlib seaborn scipy scikit-learn tenacity

# Standard library imports
import os
import sys
import datetime
import logging
from collections import deque
import random
import time

# Third-party library imports
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
from sklearn.ensemble import IsolationForest
from tenacity import retry, wait_exponential, stop_after_attempt, before_log

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def create_initial_report_state(
    start_date: datetime.date,
    end_date: datetime.date,
    report_name: str = "Monthly Compliance & Audit Report"
) -> dict:
    """
    Initializes the state dictionary for the monthly performance report.

    Parameters
    ----------
    start_date : datetime.date
        The start date for the report period.
    end_date : datetime.date
        The end date for the report period.
    report_name : str, optional
        The name of the report, defaults to "Monthly Compliance & Audit Report".

    Returns
    -------
    dict
        An initialized state dictionary with report parameters and empty data structures.

    Examples
    --------
    >>> from datetime import date
    >>> state = create_initial_report_state(date(2023, 1, 1), date(2023, 1, 31))
    >>> state['report_name']
    'Monthly Compliance & Audit Report'
    >>> state['report_period_start']
    date(2023, 1, 1)
    """
    logger.info(f"Initializing report state for '{report_name}' from {start_date} to {end_date}")

    state = {
        "report_name": report_name,
        "report_period_start": start_date,
        "report_period_end": end_date,
        "data": {},
        "kpis": {},
        "anomalies": {},
        "visualizations": {},
        "logs": deque(maxlen=100) # Store recent logs
    }
    logger.debug("Initial report state created successfully.")
    return state

def simulate_compliance_data(state: dict, num_entries: int = 1000) -> dict:
    """
    Simulates synthetic compliance and audit data.

    Parameters
    ----------
    state : dict
        The current state dictionary, expected to contain 'report_period_start'
        and 'report_period_end' keys.
    num_entries : int, optional
        The number of data entries to generate, defaults to 1000.

    Returns
    -------
    dict
        The updated state dictionary with the simulated data stored under
        `state['data']['raw_compliance_data']`.

    Examples
    --------
    >>> from datetime import date
    >>> initial_state = create_initial_report_state(date(2023, 1, 1), date(2023, 1, 31))
    >>> updated_state = simulate_compliance_data(initial_state, num_entries=50)
    >>> 'raw_compliance_data' in updated_state['data']
    True
    >>> len(updated_state['data']['raw_compliance_data'])
    50
    """
    logger.info(f"Simulating {num_entries} compliance data entries.")

    start_date = state['report_period_start']
    end_date = state['report_period_end']

    dates = pd.to_datetime(pd.date_range(start=start_date, end=end_date, periods=num_entries))
    transaction_ids = [f"TRX{i:05d}" for i in range(num_entries)]
    risk_scores = np.random.normal(loc=50, scale=15, size=num_entries).clip(0, 100)
    compliance_status = np.random.choice(['Compliant', 'Non-Compliant', 'Warning'], size=num_entries, p=[0.85, 0.10, 0.05])
    audit_flag = np.random.choice([True, False], size=num_entries, p=[0.02, 0.98])
    transaction_values = np.random.lognormal(mean=7, sigma=1, size=num_entries)

    df = pd.DataFrame({
        'date': dates,
        'transaction_id': transaction_ids,
        'risk_score': risk_scores,
        'compliance_status': compliance_status,
        'audit_flag': audit_flag,
        'transaction_value': transaction_values
    })

    state['data']['raw_compliance_data'] = df
    state['logs'].append(f"Simulated {num_entries} raw compliance data records.")
    logger.debug("Compliance data simulation completed.")
    return state

def clean_compliance_data(state: dict, fill_method: str = 'mean') -> dict:
    """
    Cleans and preprocesses the raw compliance data.

    Parameters
    ----------
    state : dict
        The current state dictionary, expected to contain
        `state['data']['raw_compliance_data']`.
    fill_method : str, optional
        Method to fill missing numerical values ('mean', 'median', 'drop'),
        defaults to 'mean'.

    Returns
    -------
    dict
        The updated state dictionary with the cleaned data stored under
        `state['data']['cleaned_compliance_data']`.

    Examples
    --------
    >>> from datetime import date
    >>> initial_state = create_initial_report_state(date(2023, 1, 1), date(2023, 1, 31))
    >>> state_with_raw_data = simulate_compliance_data(initial_state, num_entries=50)
    >>> cleaned_state = clean_compliance_data(state_with_raw_data, fill_method='median')
    >>> 'cleaned_compliance_data' in cleaned_state['data']
    True
    >>> cleaned_state['data']['cleaned_compliance_data'].isnull().sum().sum()
    0
    """
    logger.info("Starting compliance data cleaning process.")

    if 'raw_compliance_data' not in state['data']:
        logger.error("Raw compliance data not found in state. Please simulate data first.")
        state['logs'].append("Error: Raw compliance data missing for cleaning.")
        return state

    df = state['data']['raw_compliance_data'].copy()

    # Handle missing numerical values
    numerical_cols = df.select_dtypes(include=np.number).columns
    if fill_method == 'mean':
        for col in numerical_cols:
            if df[col].isnull().any():
                df[col].fillna(df[col].mean(), inplace=True)
        logger.debug(f"Filled missing numerical values using '{fill_method}' method.")
    elif fill_method == 'median':
        for col in numerical_cols:
            if df[col].isnull().any():
                df[col].fillna(df[col].median(), inplace=True)
        logger.debug(f"Filled missing numerical values using '{fill_method}' method.")
    elif fill_method == 'drop':
        initial_rows = len(df)
        df.dropna(subset=numerical_cols, inplace=True)
        logger.warning(f"Dropped {initial_rows - len(df)} rows with missing numerical values.")
    else:
        logger.warning(f"Unknown fill_method '{fill_method}'. No numerical missing value imputation performed.")

    # Handle missing categorical values (fill with 'Unknown')
    categorical_cols = df.select_dtypes(include='object').columns
    for col in categorical_cols:
        if df[col].isnull().any():
            df[col].fillna('Unknown', inplace=True)
    logger.debug("Filled missing categorical values with 'Unknown'.")

    # Ensure 'date' column is datetime and set as index
    if 'date' in df.columns:
        df['date'] = pd.to_datetime(df['date'])
        df.set_index('date', inplace=True)
        logger.debug("Converted 'date' column to datetime and set as index.")

    state['data']['cleaned_compliance_data'] = df
    state['logs'].append("Compliance data cleaning completed.")
    logger.info("Compliance data cleaning process finished.")
    return state

def calculate_kpis(state: dict) -> dict:
    """
    Calculates Key Performance Indicators (KPIs) from the cleaned compliance data.

    Parameters
    ----------
    state : dict
        The current state dictionary, expected to contain
        `state['data']['cleaned_compliance_data']`.

    Returns
    -------
    dict
        The updated state dictionary with calculated KPIs stored under `state['kpis']`.

    Examples
    --------
    >>> from datetime import date
    >>> initial_state = create_initial_report_state(date(2023, 1, 1), date(2023, 1, 31))
    >>> state_with_raw_data = simulate_compliance_data(initial_state, num_entries=100)
    >>> cleaned_state = clean_compliance_data(state_with_raw_data)
    >>> kpi_state = calculate_kpis(cleaned_state)
    >>> 'total_transactions' in kpi_state['kpis']
    True
    >>> kpi_state['kpis']['compliance_rate'] > 0.5
    True
    """
    logger.info("Calculating KPIs for the compliance report.")

    if 'cleaned_compliance_data' not in state['data']:
        logger.error("Cleaned compliance data not found in state. Please clean data first.")
        state['logs'].append("Error: Cleaned compliance data missing for KPI calculation.")
        return state

    df = state['data']['cleaned_compliance_data']

    kpis = {
        'total_transactions': len(df),
        'compliant_transactions': len(df[df['compliance_status'] == 'Compliant']),
        'non_compliant_transactions': len(df[df['compliance_status'] == 'Non-Compliant']),
        'warning_transactions': len(df[df['compliance_status'] == 'Warning']),
        'compliance_rate': (len(df[df['compliance_status'] == 'Compliant']) / len(df)) * 100 if len(df) > 0 else 0,
        'average_risk_score': df['risk_score'].mean(),
        'max_risk_score': df['risk_score'].max(),
        'min_risk_score': df['risk_score'].min(),
        'audit_flags_count': df['audit_flag'].sum(),
        'total_transaction_value': df['transaction_value'].sum()
    }

    state['kpis'] = kpis
    state['logs'].append("KPIs calculated successfully.")
    logger.info("KPI calculation completed.")
    return state

def detect_anomalies(state: dict, contamination: float = 0.05) -> dict:
    """
    Detects anomalies in the cleaned compliance data using Isolation Forest.

    Parameters
    ----------
    state : dict
        The current state dictionary, expected to contain
        `state['data']['cleaned_compliance_data']`.
    contamination : float, optional
        The proportion of outliers in the data set, defaults to 0.05.

    Returns
    -------
    dict
        The updated state dictionary with an 'is_anomaly' column added to the
        DataFrame and a DataFrame of detected anomalies.

    Examples
    --------
    >>> from datetime import date
    >>> initial_state = create_initial_report_state(date(2023, 1, 1), date(2023, 1, 31))
    >>> state_with_raw_data = simulate_compliance_data(initial_state, num_entries=200)
    >>> cleaned_state = clean_compliance_data(state_with_raw_data)
    >>> anomaly_state = detect_anomalies(cleaned_state, contamination=0.03)
    >>> 'is_anomaly' in anomaly_state['data']['cleaned_compliance_data'].columns
    True
    >>> 'detected' in anomaly_state['anomalies']
    True
    """
    logger.info(f"Detecting anomalies using Isolation Forest with contamination={contamination}.")

    if 'cleaned_compliance_data' not in state['data']:
        logger.error("Cleaned compliance data not found in state. Cannot detect anomalies.")
        state['logs'].append("Error: Cleaned compliance data missing for anomaly detection.")
        return state

    df = state['data']['cleaned_compliance_data'].copy()

    # Select numerical features for anomaly detection
    features = ['risk_score', 'transaction_value']

    if not all(col in df.columns for col in features):
        logger.warning(f"Missing some features {features} for anomaly detection. Skipping.")
        state['logs'].append(f"Warning: Missing features {features} for anomaly detection.")
        # Add a default 'is_anomaly' column of False if features are missing
        df['is_anomaly'] = False
        state['data']['cleaned_compliance_data'] = df
        state['anomalies']['detected'] = pd.DataFrame(columns=df.columns)
        return state

    # Initialize and train Isolation Forest model
    model = IsolationForest(contamination=contamination, random_state=42)
    df['anomaly_score'] = model.fit_predict(df[features])

    # -1 for outliers, 1 for inliers
    df['is_anomaly'] = df['anomaly_score'] == -1

    anomalies_df = df[df['is_anomaly'] == True]

    state['data']['cleaned_compliance_data'] = df
    state['anomalies']['detected'] = anomalies_df
    state['logs'].append(f"Detected {len(anomalies_df)} anomalies.")
    logger.info(f"Anomaly detection completed. Found {len(anomalies_df)} anomalies.")
    return state

def generate_summary_report(state: dict) -> dict:
    """
    Generates a summary report from the calculated KPIs and detected anomalies.

    Parameters
    ----------
    state : dict
        The current state dictionary, expected to contain `state['kpis']` and
        `state['anomalies']['detected']`.

    Returns
    -------
    dict
        The updated state dictionary with the summary report stored under
        `state['report_summary']`.

    Examples
    --------
    >>> from datetime import date
    >>> initial_state = create_initial_report_state(date(2023, 1, 1), date(2023, 1, 31))
    >>> state_with_raw_data = simulate_compliance_data(initial_state, num_entries=100)
    >>> cleaned_state = clean_compliance_data(state_with_raw_data)
    >>> kpi_state = calculate_kpis(cleaned_state)
    >>> anomaly_state = detect_anomalies(kpi_state)
    >>> final_state = generate_summary_report(anomaly_state)
    >>> 'report_summary' in final_state
    True
    >>> final_state['report_summary']['compliance_rate'] == final_state['kpis']['compliance_rate']
    True
    """
    logger.info("Generating summary report.")

    summary = {
        "report_period_start": state['report_period_start'].isoformat(),
        "report_period_end": state['report_period_end'].isoformat(),
        "report_name": state['report_name'],
        "total_transactions": state['kpis'].get('total_transactions', 0),
        "compliant_transactions": state['kpis'].get('compliant_transactions', 0),
        "non_compliant_transactions": state['kpis'].get('non_compliant_transactions', 0),
        "warning_transactions": state['kpis'].get('warning_transactions', 0),
        "compliance_rate": f"{state['kpis'].get('compliance_rate', 0):.2f}%",
        "average_risk_score": f"{state['kpis'].get('average_risk_score', 0):.2f}",
        "max_risk_score": f"{state['kpis'].get('max_risk_score', 0):.2f}",
        "audit_flags_count": int(state['kpis'].get('audit_flags_count', 0)),
        "total_anomalies_detected": len(state['anomalies'].get('detected', pd.DataFrame())),
        "most_recent_log": state['logs'][-1] if state['logs'] else "No logs available."
    }

    state['report_summary'] = summary
    state['logs'].append("Summary report generated.")
    logger.info("Summary report generation completed.")
    return state

# 1. Initialize Report State
from datetime import date

report_start_date = date(2023, 1, 1)
report_end_date = date(2023, 1, 31)
initial_state = create_initial_report_state(report_start_date, report_end_date, "January 2023 Compliance Report")
print(f"Initial state created for report: {initial_state['report_name']} ({initial_state['report_period_start']} to {initial_state['report_period_end']})\n")

# 2. Simulate Compliance Data
state_with_raw_data = simulate_compliance_data(initial_state, num_entries=5000)
print(f"Simulated {len(state_with_raw_data['data']['raw_compliance_data'])} raw data entries.\n")

# 3. Clean Compliance Data
cleaned_state = clean_compliance_data(state_with_raw_data, fill_method='mean')
print(f"Cleaned data contains {len(cleaned_state['data']['cleaned_compliance_data'])} entries.\n")

# 4. Calculate KPIs
kpi_state = calculate_kpis(cleaned_state)
print("Calculated KPIs:\n")
for k, v in kpi_state['kpis'].items():
    print(f"  {k}: {v}")
print("\n")

# 5. Detect Anomalies
anomaly_state = detect_anomalies(kpi_state, contamination=0.01)
num_anomalies = len(anomaly_state['anomalies']['detected'])
print(f"Detected {num_anomalies} anomalies.\n")

# 6. Generate Summary Report
final_report_state = generate_summary_report(anomaly_state)

print("--- Final Compliance & Audit Summary Report ---")
for key, value in final_report_state['report_summary'].items():
    print(f"{key.replace('_', ' ').title()}: {value}")
print("-----------------------------------------------")

print("\nMost recent logs:")
for log_entry in list(final_report_state['logs']): # Convert deque to list for printing
    print(f"- {log_entry}")

# 1. Compliance Status Distribution
plt.figure(figsize=(8, 6))
final_report_state['data']['cleaned_compliance_data']['compliance_status'].value_counts().plot.pie(
    autopct='%1.1f%%', startangle=90, colors=sns.color_palette('pastel')
)
plt.title('Compliance Status Distribution')
plt.ylabel('') # Hide the default 'compliance_status' label
plt.show()

# 2. Risk Score Distribution
plt.figure(figsize=(10, 6))
sns.histplot(final_report_state['data']['cleaned_compliance_data']['risk_score'], bins=30, kde=True)
plt.title('Distribution of Risk Scores')
plt.xlabel('Risk Score')
plt.ylabel('Number of Transactions')
plt.grid(axis='y', alpha=0.75)
plt.show()

# 3. Anomalies Over Time
anomalies_df = final_report_state['anomalies']['detected']
if not anomalies_df.empty:
    plt.figure(figsize=(12, 6))
    anomalies_df.resample('D').size().plot(marker='o', linestyle='-', color='r')
    plt.title('Number of Anomalies Detected Over Time')
    plt.xlabel('Date')
    plt.ylabel('Number of Anomalies')
    plt.grid(True)
    plt.tight_layout()
    plt.show()
else:
    print("No anomalies detected to visualize over time.")

final_report_state['logs'].append("Visualizations generated.")
logger.info("Visualizations generation completed.")