import csv
import random
from datetime import datetime, timedelta


CITIES = [
    "New York",
    "Los Angeles",
    "Chicago",
    "Toronto",
    "Mexico City",
    "Vancouver",
    "Miami",
    "Montreal",
    "San Francisco",
    "Las Vegas",
    "Houston",
    "Atlanta",
    "Seattle",
    "Boston",
    "Calgary"
]


CITY_WEATHER_CONDITIONS = {
    "New York": ["sunny", "rain", "snow", "partly cloudy", "cloudy", "thunder showers"],
    "Los Angeles": ["sunny", "partly cloudy", "fog"],
    "Chicago": ["sunny", "rain", "snow", "cloudy", "partly cloudy", "thunder showers", "freezing rain"],
    "Toronto": ["sunny", "rain", "snow", "cloudy", "partly cloudy", "thunder showers", "freezing rain"],
    "Mexico City": ["sunny", "rain", "partly cloudy", "cloudy", "thunder showers"],
    "Vancouver": ["rain", "cloudy", "fog", "partly cloudy"],
    "Miami": ["sunny", "rain", "thunder showers", "partly cloudy"],
    "Montreal": ["sunny", "rain", "snow", "cloudy", "partly cloudy", "thunder showers", "freezing rain"],
    "San Francisco": ["sunny", "fog", "partly cloudy", "cloudy"],
    "Las Vegas": ["sunny", "partly cloudy", "thunder showers"],
    "Houston": ["sunny", "rain", "partly cloudy", "thunder showers"],
    "Atlanta": ["sunny", "rain", "partly cloudy", "cloudy", "thunder showers"],
    "Seattle": ["rain", "cloudy", "fog", "partly cloudy"],
    "Boston": ["sunny", "rain", "snow", "cloudy", "partly cloudy", "thunder showers"],
    "Calgary": ["sunny", "snow", "partly cloudy", "cloudy", "thunder showers", "freezing rain"]
}


WEATHER_CONDITIONS = {
    "sunny": {"temp_range": (15, 40), "humidity_range": (10, 50), "wind_speed_range": (0, 20)},
    "cloudy": {"temp_range": (10, 25), "humidity_range": (50, 90), "wind_speed_range": (10, 30)},
    "rain": {"temp_range": (5, 20), "humidity_range": (70, 100), "wind_speed_range": (10, 25)},
    "fog": {"temp_range": (0, 15), "humidity_range": (80, 100), "wind_speed_range": (0, 10)},
    "partly cloudy": {"temp_range": (10, 30), "humidity_range": (40, 75), "wind_speed_range": (0, 20)},
    "thunder showers": {"temp_range": (18, 28), "humidity_range": (75, 100), "wind_speed_range": (15, 30)},
    "snow": {"temp_range": (-30, 0), "humidity_range": (80, 100), "wind_speed_range": (5, 15)},
    "freezing rain": {"temp_range": (-5, 3), "humidity_range": (80, 100), "wind_speed_range": (0, 10)}
}


def generate_historic_data(start_date="2018-01-01", num_days=2000):
    weather_data = {}
    current_date = datetime.strptime(start_date, "%Y-%m-%d")

    for city in CITIES:
        city_data = {}
        for _ in range(num_days):
            date_str = current_date.strftime("%Y-%m-%d")

            condition = random.choice(CITY_WEATHER_CONDITIONS[city])
            temperature = random.randint(*WEATHER_CONDITIONS[condition]["temp_range"])
            humidity = random.randint(*WEATHER_CONDITIONS[condition]["humidity_range"])

            city_data[date_str] = {
                "temperature": temperature,
                "condition": condition,
                "humidity": humidity
            }
            current_date += timedelta(days=1)
        weather_data[city] = city_data
        current_date = datetime.strptime(start_date, "%Y-%m-%d")

    return weather_data


def generate_forecast(days=14):
    forecasts = {}

    for city in CITIES:
        forecast = {}
        for day in range(days):
            date = datetime.now() + timedelta(days=day)
            condition = random.choice(CITY_WEATHER_CONDITIONS[city])
            temperature = random.randint(*WEATHER_CONDITIONS[condition]["temp_range"])
            humidity = random.randint(*WEATHER_CONDITIONS[condition]["humidity_range"])

            daily_weather = {
                "day": date.strftime("%A"),
                "temperature": temperature,
                "humidity": humidity,
                "condition": condition,
            }
            forecast[date.strftime("%Y-%m-%d")] = daily_weather
        forecasts[city] = forecast

    return forecasts


def generate_current_weather():
    weather = {}
    for city in CITIES:
        condition = random.choice(CITY_WEATHER_CONDITIONS[city])
        temperature = random.randint(*WEATHER_CONDITIONS[condition]["temp_range"])
        humidity = random.randint(*WEATHER_CONDITIONS[condition]["humidity_range"])
        wind_speed = random.randint(*WEATHER_CONDITIONS[condition]["wind_speed_range"])

        data = {
            "temperature": temperature,
            "humidity": humidity,
            "condition": condition,
            "wind_speed": wind_speed,
        }

        weather[city] = data
    return weather


if __name__ == "__main__":
    historic = generate_historic_data()
    forecast = generate_forecast()
    current = generate_current_weather()


    with open("../data/historic_data.csv", mode="w", newline="") as file:
        writer = csv.writer(file)

        writer.writerow(["city", "date", "condition", "humidity", "temperature"])

        for city, data in historic.items():
            for date, weather in data.items():
                row = [
                    city,
                    date,
                    weather.get("condition"),
                    weather.get("humidity"),
                    weather.get("temperature"),
                ]
                writer.writerow(row)

    with open("../data/forecast_data.csv", mode="w", newline="") as file:
        writer = csv.writer(file)

        # Write the header
        writer.writerow(["city", "date", "day", "condition", "temperature", "humidity"])

        for city, data in forecast.items():
            for date, weather in data.items():
                row = [
                    city,
                    date,
                    weather.get("day"),
                    weather.get("condition"),
                    weather.get("temperature"),
                    weather.get("humidity"),
                ]
                writer.writerow(row)

    with open("../data/current_data.csv", mode="w", newline="") as file:
        writer = csv.writer(file)
        writer.writerow(["city", "condition", "temperature", "humidity", "wind_speed"])

        for city, data in current.items():
            row = [
                city,
                data.get("condition"),
                data.get("temperature"),
                data.get("humidity"),
                data.get("wind_speed"),
            ]
            writer.writerow(row)