import pandas as pd
import numpy as np
from sklearn.ensemble import IsolationForest

# Load data
df = pd.read_csv('AstraZeneca_Cleaned_Data.csv', low_memory=False)

# Ensure proper types
df['Amount in USD'] = pd.to_numeric(df['Amount in USD'], errors='coerce')
df['Pstng Date'] = pd.to_datetime(df['Pstng Date'], errors='coerce')
df = df.dropna(subset=['Amount in USD', 'Pstng Date'])

print("Data Loaded. Starting Advanced Anomaly Detection...")

# --- 1. Statistical Outliers (Z-Score) ---
# Calculate Mean and Std Dev globally or by Category (User originally asked for by Category, let's keep that but maybe also global)
# Implementing Global Z-Score for broad outliers as well
df['Z_Score'] = (df['Amount in USD'] - df['Amount in USD'].mean()) / df['Amount in USD'].std()

# --- 2. Isolation Forest (Machine Learning) ---
print("Training Isolation Forest...")
iso_forest = IsolationForest(contamination=0.01, random_state=42) # Assuming 1% anomalies
df['Iso_Score'] = iso_forest.fit_predict(df[['Amount in USD']])
# Iso_Score: -1 is anomaly, 1 is normal

# --- 3. Duplicate Detection ---
print("Checking for Potential Duplicates...")
# Logic: Same Amount, Same Vendor (Name of offsetting account), Same Reference (Ref.key (header) 1)
# Within 48 hours
potential_dupes = df.sort_values(['Name of offsetting account', 'Ref.key (header) 1', 'Amount in USD', 'Pstng Date'])

potential_dupes['Time_Diff'] = potential_dupes.groupby(
    ['Name of offsetting account', 'Ref.key (header) 1', 'Amount in USD']
)['Pstng Date'].diff()

# Flag if Time_Diff is present and less than 2 days (48 hours)
# Note: diff() gives NaT for the first item in the group. We only care about the subsequent ones which are "dupes" of the first.
potential_dupes['Is_Duplicate'] = potential_dupes['Time_Diff'] <= pd.Timedelta(hours=48)

# Merge duplicate flag back to main df
df['Is_Duplicate'] = False
df.loc[potential_dupes[potential_dupes['Is_Duplicate']].index, 'Is_Duplicate'] = True


# --- 4. Triage / Risk Consolidation ---
print("Consolidating Risk Report...")

def determine_risk(row):
    risks = []
    
    # 1. High Z-Score (Statistical Outlier)
    if abs(row['Z_Score']) > 3:
        risks.append("high_z_score")
    
    # 2. Isolation Forest Anomaly
    if row['Iso_Score'] == -1:
        risks.append("ml_anomaly")
        
    # 3. Duplicate
    if row['Is_Duplicate']:
        risks.append("potential_duplicate")
        
    if not risks:
        return "Normal"
    else:
        return ", ".join(risks)

df['Risk_Type'] = df.apply(determine_risk, axis=1)

# Filter for anything that is NOT Normal
audit_risks = df[df['Risk_Type'] != 'Normal'].copy()

# Add a "Priority" column
# Duplicates = High, Both Z-Score and ML = High, others Medium
def assign_priority(risk_type):
    if "potential_duplicate" in risk_type:
        return "High"
    if "high_z_score" in risk_type and "ml_anomaly" in risk_type:
        return "High"
    return "Medium"

audit_risks['Priority'] = audit_risks['Risk_Type'].apply(assign_priority)

# Select relevant columns for the Audit File
audit_columns = ['Pstng Date', 'DocumentNo', 'Clrng doc.', 'Name of offsetting account', 
                 'Amount in USD', 'Category', 'Z_Score', 'Iso_Score', 'Risk_Type', 'Priority']
# Ensure these columns exist
existing_cols = [c for c in audit_columns if c in audit_risks.columns]
final_audit = audit_risks[existing_cols]

try:
    final_audit.to_csv('Audit_Risks.csv', index=False)
    final_audit.to_excel('Audit_Risks.xlsx', index=False)
    print(f"Analysis Complete. Generated 'Audit_Risks.csv' and 'Audit_Risks.xlsx' with {len(final_audit)} flagged transactions.")
    print(final_audit['Risk_Type'].value_counts())
except PermissionError:
    print("\n[ERROR] Could not save the file. Please close 'Audit_Risks.csv' or 'Audit_Risks.xlsx' and try again.")
except Exception as e:
    print(f"\n[ERROR] An unexpected error occurred while saving: {e}")