from numpy import array, arange, random, where, clip, zeros
import os
import csv

def load_or_generate_data(num_days=30, output_path="pjm_data.csv", seed=42):
    """
    Simulates fetching data from PJM DataMiner. Now uses ONLY numpy and csv
    to minimize Docker build time and image size (removes Pandas).
    """
    if output_path and os.path.exists(output_path):
        # Load using numpy's structured array or dict of arrays
        data = {}
        with open(output_path, 'r') as f:
            reader = csv.DictReader(f)
            rows = list(reader)
            for key in rows[0].keys():
                data[key] = array([float(r[key]) for r in rows])
        return data
        
    random.seed(seed)
    total_hours = num_days * 24
    hours = arange(total_hours)
    hours_of_day = hours % 24
    
    # 1. Real-Time LMP ($/MWh) - Diurnal pattern
    base_price = random.normal(30, 5, total_hours)
    peak_multiplier = where((hours_of_day >= 16) & (hours_of_day <= 20), 2.5, 1.0)
    lmp = base_price * peak_multiplier + random.normal(0, 10, total_hours)
    lmp = clip(lmp, 10, 300)
    
    # 2. Hourly Load (MW) - Peak Shaving calibration
    base_load = random.normal(15, 2, total_hours)
    load_multiplier = where((hours_of_day >= 9) & (hours_of_day <= 18), 1.5, 1.0)
    load = base_load * load_multiplier + random.normal(0, 1.0, total_hours)
    load = clip(load, 5, 50)
    
    # 3. RegD Signal (FR signal tracking)
    regd = zeros(total_hours)
    theta, mu, sigma = 0.15, 0.0, 0.2
    for i in range(1, total_hours):
        regd[i] = regd[i-1] + theta * (mu - regd[i-1]) + sigma * random.normal()
    regd = clip(regd, -1.0, 1.0)
    
    data = {
        "hour_of_day": hours_of_day,
        "lmp": lmp,
        "load": load,
        "regd": regd
    }
    
    if output_path:
        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
        keys = data.keys()
        with open(output_path, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(keys)
            writer.writerows(zip(*[data[k] for k in keys]))
        
    return data