import pandas as pd
import os

# Create data directory if it doesn't exist
os.makedirs('data', exist_ok=True)

# TechHealth AI data
company_data = {
    "name": "TechHealth AI",
    "stage": "Seed",
    "founded": "18 months ago",
    "employees": 12,
    "last_funding": "$1.2M seed round 10 months ago",
    "cash": 320000,
    "burn_rate": 85000,
    "revenue": 15000,
    "growth_rate": 0.08
}

# Save company data
pd.DataFrame([company_data]).to_csv('data/startup_data.csv', index=False)

# Cash flow history
cash_flow_data = {
    "Month": [f"Month {i}" for i in range(1, 11)],
    "Revenue": [8000, 8500, 9200, 10000, 10800, 11700, 12600, 13600, 14700, 15800],
    "Payroll": [60000, 60000, 62000, 62000, 65000, 65000, 70000, 70000, 75000, 75000],
    "Marketing": [8000, 9000, 10000, 12000, 15000, 18000, 15000, 12000, 10000, 8000],
    "Office": [5000, 5000, 5000, 5000, 5000, 5000, 5000, 5000, 5000, 5000],
    "Software": [3000, 3200, 3500, 3800, 4000, 4200, 4500, 4800, 5000, 5200],
    "Travel": [2000, 1800, 2500, 3000, 4000, 4500, 3500, 3000, 2500, 2000],
    "Legal": [1500, 1000, 800, 1200, 800, 2000, 1500, 1000, 3000, 1200],
    "Misc": [1000, 1200, 1300, 1500, 1700, 1800, 2000, 2200, 2500, 2800]
}

# Add calculated fields
df = pd.DataFrame(cash_flow_data)
df["Total_Expenses"] = df[["Payroll", "Marketing", "Office", "Software", "Travel", "Legal", "Misc"]].sum(axis=1)
df["Net_Burn"] = df["Total_Expenses"] - df["Revenue"]

# Save cash flow data
df.to_csv('data/projections.csv', index=False)

# Transaction data
transactions = pd.DataFrame([
    {"Date": "2023-11-05", "Category": "Travel", "Vendor": "Caribbean Cruises", "Amount": 8500, "Description": "Team Retreat Planning", "Flag": "Suspicious"},
    {"Date": "2023-11-12", "Category": "Marketing", "Vendor": "LuxuryGifts Inc", "Amount": 4200, "Description": "Client Appreciation", "Flag": "Suspicious"},
    {"Date": "2023-11-22", "Category": "Office", "Vendor": "Premium Furniture", "Amount": 12000, "Description": "Office Upgrades", "Flag": "Suspicious"},
    {"Date": "2023-11-28", "Category": "Consulting", "Vendor": "Strategic Vision LLC", "Amount": 7500, "Description": "Strategy Consulting", "Flag": "Suspicious"},
    {"Date": "2023-12-05", "Category": "Software", "Vendor": "Personal Apple Store", "Amount": 3200, "Description": "Development Tools", "Flag": "Suspicious"},
    {"Date": "2023-12-12", "Category": "Legal", "Vendor": "Anderson Brothers", "Amount": 5800, "Description": "Legal Services", "Flag": "Normal"},
    {"Date": "2023-12-20", "Category": "Payroll", "Vendor": "November Payroll", "Amount": 75000, "Description": "Monthly Payroll", "Flag": "Normal"},
    {"Date": "2023-12-22", "Category": "Marketing", "Vendor": "Google Ads", "Amount": 8000, "Description": "Ad Campaign", "Flag": "Normal"},
    {"Date": "2023-12-25", "Category": "Office", "Vendor": "WeWork", "Amount": 5000, "Description": "Monthly Rent", "Flag": "Normal"},
    {"Date": "2023-12-28", "Category": "Software", "Vendor": "AWS", "Amount": 5200, "Description": "Cloud Services", "Flag": "Normal"},
    {"Date": "2024-01-05", "Category": "Travel", "Vendor": "Delta Airlines", "Amount": 1200, "Description": "Client Meeting Travel", "Flag": "Normal"},
    {"Date": "2024-01-10", "Category": "Marketing", "Vendor": "Facebook Ads", "Amount": 4500, "Description": "Social Media Campaign", "Flag": "Normal"},
    {"Date": "2024-01-15", "Category": "Software", "Vendor": "Atlassian", "Amount": 2800, "Description": "Development Tools", "Flag": "Normal"},
    {"Date": "2024-01-20", "Category": "Payroll", "Vendor": "January Payroll", "Amount": 75000, "Description": "Monthly Payroll", "Flag": "Normal"},
    {"Date": "2024-01-25", "Category": "Office", "Vendor": "WeWork", "Amount": 5000, "Description": "Monthly Rent", "Flag": "Normal"}
])

# Save transactions data
transactions.to_csv('data/transactions.csv', index=False)

# Create a separate file for suspicious transactions for easier analysis
suspicious = transactions[transactions['Flag'] == 'Suspicious']
suspicious.to_csv('data/suspicious.csv', index=False)

print("Sample data files have been created in the 'data' directory.")