"""
================================================================================
DATATHON 2025 - ASTRAZENECA CASH FLOW CHALLENGE
01_data_cleaning.py - Data Processing Pipeline
================================================================================

1. Load and clean main transaction data
2. Remove duplicate transactions
3. Join supporting tables (Category Linkage, Country Mapping, Exchange Rate)
4. Create weekly aggregation for forecasting
5. Create train/test split
6. Generate data quality report


"""

import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import warnings
warnings.filterwarnings('ignore')

# ================================================================================
# CONFIGURATION
# ================================================================================

# File paths
DATA_PATH = "dataset/"
MAIN_FILE = DATA_PATH + "Datathon Dataset.xlsx - Data - Main.csv"
CASH_BALANCE_FILE = DATA_PATH + "Datathon Dataset.xlsx - Data - Cash Balance.csv"
CATEGORY_FILE = DATA_PATH + "Datathon Dataset.xlsx - Others - Category Linkage.csv"
COUNTRY_FILE = DATA_PATH + "Datathon Dataset.xlsx - Others - Country Mapping.csv"
EXCHANGE_FILE = DATA_PATH + "Datathon Dataset.xlsx - Others - Exchange Rate.csv"

# Output files
OUTPUT_CLEANED = "data_cleaned_main.csv"
OUTPUT_WEEKLY = "data_cleaned_weekly.csv"
OUTPUT_TRAIN = "data_train_weekly.csv"
OUTPUT_TEST = "data_test_weekly.csv"
OUTPUT_REPORT = "DATA_QUALITY_REPORT.txt"

# Train/Test split configuration
TRAIN_RATIO = 0.8  # 80% train, 20% test

print("=" * 80)
print("DATATHON 2025 - DATA PROCESSING PIPELINE")
print("=" * 80)

# ================================================================================
# STEP 1: LOAD RAW DATA
# ================================================================================

print("\n[STEP 1] Loading raw data files...")

# Load main transaction data
df_main = pd.read_csv(MAIN_FILE)
print(f"  - Main transactions loaded: {len(df_main):,} rows")

# Load supporting tables
df_category = pd.read_csv(CATEGORY_FILE)
print(f"  - Category Linkage loaded: {len(df_category):,} rows")

df_country = pd.read_csv(COUNTRY_FILE)
print(f"  - Country Mapping loaded: {len(df_country):,} rows")

df_exchange = pd.read_csv(EXCHANGE_FILE)
print(f"  - Exchange Rates loaded: {len(df_exchange):,} rows")

df_cash_balance = pd.read_csv(CASH_BALANCE_FILE)
print(f"  - Cash Balance loaded: {len(df_cash_balance):,} rows")

# ================================================================================
# STEP 2: DATA CLEANING & PARSING
# ================================================================================

print("\n[STEP 2] Cleaning and parsing data...")

# Parse date columns
# Pstng Date is the primary date field for cash flow timing
df_main['Pstng Date'] = pd.to_datetime(df_main['Pstng Date'], format='%m/%d/%Y', errors='coerce')

# Parse Doc..Date (original document date)
df_main['Doc..Date'] = pd.to_datetime(df_main['Doc..Date'], format='%m/%d/%y %H:%M', errors='coerce')

# Parse File Month
df_main['File Month'] = pd.to_datetime(df_main['File Month'], format='%A, %d %B, %Y', errors='coerce')

print(f"  - Date columns parsed")

# Convert Amount in USD to numeric (handle comma separators)
df_main['Amount in USD'] = df_main['Amount in USD'].astype(str).str.replace(',', '').astype(float)
df_main['Amt in loc.cur.'] = df_main['Amt in loc.cur.'].astype(str).str.replace(',', '').astype(float)
print(f"  - Amount columns converted to numeric (comma separators handled)")

# Check for parsing errors
date_nulls = df_main['Pstng Date'].isnull().sum()
if date_nulls > 0:
    print(f"  - WARNING: {date_nulls} rows with invalid Pstng Date")

amount_nulls = df_main['Amount in USD'].isnull().sum()
if amount_nulls > 0:
    print(f"  - WARNING: {amount_nulls} rows with null Amount in USD")

# ================================================================================
# STEP 3: IDENTIFY AND FLAG POTENTIAL DUPLICATES (NOT REMOVED)
# ================================================================================

print("\n[STEP 3] Analyzing potential duplicates...")

original_count = len(df_main)

# Note: DocumentNo is NOT globally unique - same DocumentNo can exist across entities
# This is normal for SAP systems where DocumentNo is entity-specific
duplicate_doc_count = df_main['DocumentNo'].duplicated().sum()
print(f"  - DocumentNo entries shared across entities: {duplicate_doc_count:,} (normal for SAP)")

# Check for TRUE duplicates (same entity + same date + same amount + same DocumentNo)
# These would indicate actual duplicate entries that need investigation
true_duplicate_cols = ['DocumentNo', 'Pstng Date', 'Amount in USD', 'Name', 'Category']
df_main['potential_duplicate'] = df_main.duplicated(subset=true_duplicate_cols, keep=False)
potential_duplicates = df_main['potential_duplicate'].sum()
print(f"  - Potential true duplicates (for investigation): {potential_duplicates:,}")

# Flag and save potential duplicates for anomaly detection
if potential_duplicates > 0:
    df_flagged = df_main[df_main['potential_duplicate']].copy()
    df_flagged.to_csv("data_flagged_potential_duplicates.csv", index=False)
    print(f"  - Flagged for review: data_flagged_potential_duplicates.csv")

# We do NOT remove duplicates automatically - finance team should review
# All transactions are retained with flags
df_clean = df_main.drop(columns=['potential_duplicate']).copy()
print(f"  - Clean dataset: {len(df_clean):,} rows (all retained, duplicates flagged only)")

# ================================================================================
# STEP 4: JOIN SUPPORTING TABLES
# ================================================================================

print("\n[STEP 4] Joining supporting tables...")

# First, normalize category names in main data to match Category Linkage
# Fix common naming inconsistencies
category_mapping = {
    'Non Netting AP': 'Non-Netting AP',
    'Non Netting AR': 'Non-Netting AR', 
    'Dividend payout': 'Dividend Payout',
}
df_clean['Category_Normalized'] = df_clean['Category'].replace(category_mapping)
print(f"  - Category names normalized ({len(category_mapping)} mappings applied)")

# Join Category Linkage
df_category_clean = df_category.rename(columns={
    'Category Names': 'Category_Name',
    'Category': 'Flow_Direction',
    'Cat Order': 'Flow_Order'
})

df_clean = df_clean.merge(
    df_category_clean[['Category_Name', 'Flow_Direction', 'Flow_Order']],
    left_on='Category_Normalized',
    right_on='Category_Name',
    how='left'
)

# For 'Other' category (not in linkage), classify based on amount sign
other_mask = df_clean['Category'] == 'Other'
df_clean.loc[other_mask & (df_clean['Amount in USD'] > 0), 'Flow_Direction'] = 'Inflow'
df_clean.loc[other_mask & (df_clean['Amount in USD'] <= 0), 'Flow_Direction'] = 'Outflow'
df_clean.loc[other_mask & (df_clean['Amount in USD'] > 0), 'Flow_Order'] = 2
df_clean.loc[other_mask & (df_clean['Amount in USD'] <= 0), 'Flow_Order'] = 1
print(f"  - Category Linkage joined (Inflow/Outflow classification)")

# Join Country Mapping
df_clean = df_clean.merge(
    df_country[['Code', 'Country']],
    left_on='Name',
    right_on='Code',
    how='left'
)
print(f"  - Country Mapping joined")

# Clean up duplicate/temporary columns
if 'Code' in df_clean.columns:
    df_clean = df_clean.drop(columns=['Code'])
if 'Category_Name' in df_clean.columns:
    df_clean = df_clean.drop(columns=['Category_Name'])
if 'Category_Normalized' in df_clean.columns:
    df_clean = df_clean.drop(columns=['Category_Normalized'])

# Verify join success
unmatched_categories = df_clean['Flow_Direction'].isnull().sum()
unmatched_countries = df_clean['Country'].isnull().sum()
print(f"  - Unmatched categories: {unmatched_categories}")
print(f"  - Unmatched countries: {unmatched_countries}")

# ================================================================================
# STEP 5: CREATE DERIVED FIELDS
# ================================================================================

print("\n[STEP 5] Creating derived fields...")

# Week number and year for weekly aggregation
df_clean['Year'] = df_clean['Pstng Date'].dt.year
df_clean['Week'] = df_clean['Pstng Date'].dt.isocalendar().week
df_clean['Year_Week'] = df_clean['Pstng Date'].dt.strftime('%Y-W%V')
df_clean['Week_Start'] = df_clean['Pstng Date'] - pd.to_timedelta(df_clean['Pstng Date'].dt.dayofweek, unit='D')
df_clean['Week_Start'] = df_clean['Week_Start'].dt.normalize()

# Month for monthly analysis
df_clean['Year_Month'] = df_clean['Pstng Date'].dt.strftime('%Y-%m')

# ============================================================================
# CASH FLOW SIGN CONVENTION (IMPORTANT!)
# ============================================================================
# Positive amounts (+) = INFLOW  = Money coming IN  (received from customers, loans, etc.)
# Negative amounts (-) = OUTFLOW = Money going OUT  (paid to vendors, payroll, taxes, etc.)
# ============================================================================
# Example:
#   +$100,000 = Company received $100K (e.g., customer payment, AR)
#   -$50,000  = Company paid out $50K (e.g., vendor payment, AP, payroll)
#   Net Cash Flow = Inflows + Outflows (outflows are negative, so effectively subtracted)
# ============================================================================

df_clean['Is_Inflow'] = df_clean['Amount in USD'] > 0
df_clean['Is_Outflow'] = df_clean['Amount in USD'] < 0

print(f"  - Week fields created (Year_Week, Week_Start)")
print(f"  - Inflow/Outflow flags created")

# ================================================================================
# STEP 6: SAVE CLEANED TRANSACTION DATA
# ================================================================================

print("\n[STEP 6] Saving cleaned transaction data...")

df_clean.to_csv(OUTPUT_CLEANED, index=False)
print(f"  - Saved: {OUTPUT_CLEANED} ({len(df_clean):,} rows)")

# ================================================================================
# STEP 7: CREATE WEEKLY AGGREGATION
# ================================================================================

print("\n[STEP 7] Creating weekly aggregation...")

# Aggregate by week for forecasting
weekly_agg = df_clean.groupby(['Week_Start']).agg({
    'Amount in USD': 'sum',
    'DocumentNo': 'count'
}).reset_index()

weekly_agg.columns = ['Week_Start', 'Net_Cash_Flow_USD', 'Transaction_Count']

# Add week identifiers
weekly_agg['Week_Start'] = pd.to_datetime(weekly_agg['Week_Start'])
weekly_agg = weekly_agg.sort_values('Week_Start').reset_index(drop=True)
weekly_agg['Week_Number'] = range(1, len(weekly_agg) + 1)

# Calculate cumulative cash flow
weekly_agg['Cumulative_Cash_Flow'] = weekly_agg['Net_Cash_Flow_USD'].cumsum()

# Calculate inflows and outflows separately
weekly_inflow = df_clean[df_clean['Is_Inflow']].groupby('Week_Start')['Amount in USD'].sum().reset_index()
weekly_inflow.columns = ['Week_Start', 'Total_Inflow_USD']

weekly_outflow = df_clean[df_clean['Is_Outflow']].groupby('Week_Start')['Amount in USD'].sum().reset_index()
weekly_outflow.columns = ['Week_Start', 'Total_Outflow_USD']

weekly_agg = weekly_agg.merge(weekly_inflow, on='Week_Start', how='left')
weekly_agg = weekly_agg.merge(weekly_outflow, on='Week_Start', how='left')

# Fill NaN with 0
weekly_agg['Total_Inflow_USD'] = weekly_agg['Total_Inflow_USD'].fillna(0)
weekly_agg['Total_Outflow_USD'] = weekly_agg['Total_Outflow_USD'].fillna(0)

# Add entity-level breakdown
for entity in df_clean['Name'].unique():
    entity_weekly = df_clean[df_clean['Name'] == entity].groupby('Week_Start')['Amount in USD'].sum().reset_index()
    entity_weekly.columns = ['Week_Start', f'Net_Cash_{entity}']
    weekly_agg = weekly_agg.merge(entity_weekly, on='Week_Start', how='left')
    weekly_agg[f'Net_Cash_{entity}'] = weekly_agg[f'Net_Cash_{entity}'].fillna(0)

# Add category-level breakdown for top categories
top_categories = df_clean.groupby('Category')['Amount in USD'].apply(lambda x: abs(x).sum()).nlargest(10).index.tolist()
for category in top_categories:
    cat_weekly = df_clean[df_clean['Category'] == category].groupby('Week_Start')['Amount in USD'].sum().reset_index()
    cat_weekly.columns = ['Week_Start', f'Net_Cash_{category.replace(" ", "_")}']
    weekly_agg = weekly_agg.merge(cat_weekly, on='Week_Start', how='left')
    weekly_agg[f'Net_Cash_{category.replace(" ", "_")}'] = weekly_agg[f'Net_Cash_{category.replace(" ", "_")}'].fillna(0)

print(f"  - Weekly aggregation created: {len(weekly_agg)} weeks")
print(f"  - Date range: {weekly_agg['Week_Start'].min()} to {weekly_agg['Week_Start'].max()}")

# Save weekly data
weekly_agg.to_csv(OUTPUT_WEEKLY, index=False)
print(f"  - Saved: {OUTPUT_WEEKLY}")

# ================================================================================
# STEP 8: CREATE TRAIN/TEST SPLIT
# ================================================================================

print("\n[STEP 8] Creating train/test split...")

total_weeks = len(weekly_agg)
train_weeks = int(total_weeks * TRAIN_RATIO)
test_weeks = total_weeks - train_weeks

df_train = weekly_agg.iloc[:train_weeks].copy()
df_test = weekly_agg.iloc[train_weeks:].copy()

print(f"  - Total weeks: {total_weeks}")
print(f"  - Train weeks: {train_weeks} ({TRAIN_RATIO*100:.0f}%)")
print(f"  - Test weeks: {test_weeks} ({(1-TRAIN_RATIO)*100:.0f}%)")
print(f"  - Train period: {df_train['Week_Start'].min()} to {df_train['Week_Start'].max()}")
print(f"  - Test period: {df_test['Week_Start'].min()} to {df_test['Week_Start'].max()}")

# Save train/test sets
df_train.to_csv(OUTPUT_TRAIN, index=False)
df_test.to_csv(OUTPUT_TEST, index=False)
print(f"  - Saved: {OUTPUT_TRAIN}")
print(f"  - Saved: {OUTPUT_TEST}")

# ================================================================================
# STEP 9: PROCESS CASH BALANCE DATA
# ================================================================================

print("\n[STEP 9] Processing cash balance data...")

# Clean cash balance data (remove formatting)
df_cash_balance['Carryforward Balance'] = df_cash_balance['Carryforward Balance'].str.replace(',', '').str.strip().astype(float)
df_cash_balance['Carryforward Balance (USD)'] = df_cash_balance['Carryforward Balance (USD)'].str.replace(',', '').str.strip().astype(float)

# Calculate total starting balance in USD
total_starting_balance_usd = df_cash_balance['Carryforward Balance (USD)'].sum()
print(f"  - Total starting balance (USD): ${total_starting_balance_usd:,.2f}")

# Save processed cash balance
df_cash_balance.to_csv("data_cash_balance_processed.csv", index=False)
print(f"  - Saved: data_cash_balance_processed.csv")

# ================================================================================
# STEP 10: GENERATE DATA QUALITY REPORT
# ================================================================================

print("\n[STEP 10] Generating data quality report...")

report = []
report.append("=" * 80)
report.append("DATATHON 2025 - DATA QUALITY REPORT")
report.append("=" * 80)
report.append(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
report.append("")

# Dataset Overview
report.append("-" * 80)
report.append("DATASET OVERVIEW")
report.append("-" * 80)
report.append(f"Total transactions loaded: {original_count:,}")
report.append(f"Final clean transactions: {len(df_clean):,}")
report.append(f"Note: Duplicates are flagged but NOT removed (require manual review)")
report.append("")

# Date Range
report.append("-" * 80)
report.append("DATE RANGE")
report.append("-" * 80)
report.append(f"Earliest transaction: {df_clean['Pstng Date'].min()}")
report.append(f"Latest transaction: {df_clean['Pstng Date'].max()}")
report.append(f"Total weeks: {total_weeks}")
report.append("")

# Entity Summary
report.append("-" * 80)
report.append("ENTITY SUMMARY")
report.append("-" * 80)
entity_summary = df_clean.groupby('Name').agg({
    'Amount in USD': ['sum', 'count', 'mean'],
    'Country': 'first'
}).reset_index()
entity_summary.columns = ['Entity', 'Net_Cash_USD', 'Transaction_Count', 'Avg_Transaction', 'Country']

for _, row in entity_summary.iterrows():
    report.append(f"{row['Entity']} ({row['Country']}): ${row['Net_Cash_USD']:,.2f} ({row['Transaction_Count']:,} transactions)")
report.append("")

# Category Summary
report.append("-" * 80)
report.append("CATEGORY SUMMARY (Top 10 by Volume)")
report.append("-" * 80)
category_summary = df_clean.groupby('Category').agg({
    'Amount in USD': ['sum', 'count'],
    'Flow_Direction': 'first'
}).reset_index()
category_summary.columns = ['Category', 'Net_Cash_USD', 'Transaction_Count', 'Flow_Direction']
category_summary['Abs_Volume'] = category_summary['Net_Cash_USD'].abs()
category_summary = category_summary.sort_values('Abs_Volume', ascending=False).head(10)

for _, row in category_summary.iterrows():
    direction = row['Flow_Direction'] if pd.notna(row['Flow_Direction']) else 'Unknown'
    report.append(f"{row['Category']} ({direction}): ${row['Net_Cash_USD']:,.2f} ({row['Transaction_Count']:,} transactions)")
report.append("")

# Cash Flow Summary
report.append("-" * 80)
report.append("CASH FLOW SUMMARY")
report.append("-" * 80)
report.append("")
report.append("*** SIGN CONVENTION ***")
report.append("  Positive (+) = INFLOW  = Money received (coming IN to the company)")
report.append("  Negative (-) = OUTFLOW = Money paid out (going OUT of the company)")
report.append("")
total_inflow = df_clean[df_clean['Is_Inflow']]['Amount in USD'].sum()
total_outflow = df_clean[df_clean['Is_Outflow']]['Amount in USD'].sum()
net_cash = total_inflow + total_outflow

report.append(f"Total Inflow (USD): ${total_inflow:,.2f}")
report.append(f"Total Outflow (USD): ${total_outflow:,.2f}  <-- Negative = money going OUT")
report.append(f"Net Cash Flow (USD): ${net_cash:,.2f}")
report.append(f"Starting Balance (USD): ${total_starting_balance_usd:,.2f}")
report.append(f"Projected Ending Balance (USD): ${total_starting_balance_usd + net_cash:,.2f}")
report.append("")

# Data Quality Metrics
report.append("-" * 80)
report.append("DATA QUALITY METRICS")
report.append("-" * 80)
report.append(f"Missing Pstng Date: {df_clean['Pstng Date'].isnull().sum()}")
report.append(f"Missing Amount in USD: {df_clean['Amount in USD'].isnull().sum()}")
report.append(f"Missing Category: {df_clean['Category'].isnull().sum()}")
report.append(f"Missing Entity (Name): {df_clean['Name'].isnull().sum()}")
report.append(f"Unmatched Categories: {unmatched_categories}")
report.append(f"Unmatched Countries: {unmatched_countries}")
report.append("")

# Train/Test Split Summary
report.append("-" * 80)
report.append("TRAIN/TEST SPLIT")
report.append("-" * 80)
report.append(f"Train weeks: {train_weeks} ({TRAIN_RATIO*100:.0f}%)")
report.append(f"Train period: {df_train['Week_Start'].min()} to {df_train['Week_Start'].max()}")
report.append(f"Train net cash flow: ${df_train['Net_Cash_Flow_USD'].sum():,.2f}")
report.append("")
report.append(f"Test weeks: {test_weeks} ({(1-TRAIN_RATIO)*100:.0f}%)")
report.append(f"Test period: {df_test['Week_Start'].min()} to {df_test['Week_Start'].max()}")
report.append(f"Test net cash flow: ${df_test['Net_Cash_Flow_USD'].sum():,.2f}")
report.append("")

# Output Files
report.append("-" * 80)
report.append("OUTPUT FILES CREATED")
report.append("-" * 80)
report.append(f"1. {OUTPUT_CLEANED} - Clean transaction data ({len(df_clean):,} rows)")
report.append(f"2. {OUTPUT_WEEKLY} - Weekly aggregation ({len(weekly_agg)} weeks)")
report.append(f"3. {OUTPUT_TRAIN} - Training data ({train_weeks} weeks)")
report.append(f"4. {OUTPUT_TEST} - Test data ({test_weeks} weeks)")
report.append(f"5. data_flagged_potential_duplicates.csv - Flagged for review (if any)")
report.append(f"6. data_cash_balance_processed.csv - Starting balances (8 entities)")
report.append(f"7. {OUTPUT_REPORT} - This report")
report.append("")

report.append("=" * 80)
report.append("END OF REPORT")
report.append("=" * 80)

# Write report
report_text = "\n".join(report)
with open(OUTPUT_REPORT, 'w') as f:
    f.write(report_text)

print(f"  - Saved: {OUTPUT_REPORT}")

# Print summary
print("\n" + "=" * 80)
print("DATA PROCESSING COMPLETE")
print("=" * 80)
print(f"\nKey Metrics:")
print(f"  - Clean transactions: {len(df_clean):,}")
print(f"  - Total weeks: {total_weeks}")
print(f"  - Net cash flow: ${net_cash:,.2f}")
print(f"  - Starting balance: ${total_starting_balance_usd:,.2f}")
print(f"  - Train/Test split: {train_weeks}/{test_weeks} weeks")
print(f"\nOutput files created in project root directory.")
print("=" * 80)