from tabpfn import TabPFNTimeSeries
import pandas as pd

# Load historical sales data
sales_data = pd.read_csv("sales_timeseries.csv")
sales_data['date'] = pd.to_datetime(sales_data['date'])
sales_data = sales_data.set_index('date')

# Initialize forecasting model
ts_model = TabPFNTimeSeries()

# Forecast next quarter
quarterly_forecast = ts_model.forecast(
    data=sales_data,
    horizon=90,  # 90 days ahead
    confidence_level=0.95
)

# Business planning insights
print(f"Q4 Sales Forecast: ${quarterly_forecast.predictions.sum():,.0f}")
print(f"Confidence Range: ${quarterly_forecast.confidence_intervals[:, 0].sum():,.0f} - ${quarterly_forecast.confidence_intervals[:, 1].sum():,.0f}")

# Multi-product demand forecasting
def forecast_product_demand(product_data, horizon_days=30):
    ts_model = TabPFNTimeSeries()
    
    # Forecast demand for each product
    forecasts = {}
    for product in product_data.columns:
        product_series = product_data[product]
        forecast = ts_model.forecast(
            data=product_series,
            horizon=horizon_days,
            confidence_level=0.90
        )
        
        forecasts[product] = {
            'demand': forecast.predictions,
            'confidence_range': forecast.confidence_intervals,
            'trend': forecast.trend,
            'seasonality': forecast.seasonality
        }
    
    return forecasts

# Inventory planning
def plan_inventory(forecasts, current_stock):
    inventory_plan = {}
    
    for product, forecast in forecasts.items():
        predicted_demand = forecast['demand'].sum()
        safety_stock = forecast['confidence_range'][:, 1].sum() - predicted_demand
        
        inventory_plan[product] = {
            'recommended_order': max(0, predicted_demand - current_stock[product] + safety_stock),
            'reorder_point': safety_stock,
            'confidence_level': 0.90
        }
    
    return inventory_plan

# Financial market forecasting
def forecast_market_trends(market_data, horizon_days=30):
    ts_model = TabPFNTimeSeries()
    
    # Multi-variate forecasting for correlated markets
    market_forecast = ts_model.forecast_multivariate(
        data=market_data,
        horizon=horizon_days,
        series_names=['stock_price', 'volume', 'volatility'],
        confidence_level=0.95
    )
    
    # Risk assessment
    risk_metrics = {
        'volatility_forecast': market_forecast['volatility'].std(),
        'trend_direction': 'bullish' if market_forecast['stock_price'].iloc[-1] > market_forecast['stock_price'].iloc[0] else 'bearish',
        'confidence_level': 0.95
    }
    
    return market_forecast, risk_metrics

# Portfolio optimization
def optimize_portfolio(forecasts, current_portfolio):
    # Use forecasts to rebalance portfolio
    recommended_weights = {}
    
    for asset, forecast in forecasts.items():
        expected_return = forecast['predictions'].mean()
        risk_score = forecast['confidence_intervals'].std()
        
        # Risk-adjusted return
        risk_adjusted_return = expected_return / risk_score
        recommended_weights[asset] = risk_adjusted_return
    
    # Normalize weights
    total_weight = sum(recommended_weights.values())
    normalized_weights = {k: v/total_weight for k, v in recommended_weights.items()}
    
    return normalized_weights

# Energy consumption forecasting
def forecast_energy_consumption(consumption_data, horizon_hours=24):
    ts_model = TabPFNTimeSeries()
    
    # Hourly consumption forecast
    consumption_forecast = ts_model.forecast(
        data=consumption_data,
        horizon=horizon_hours,
        confidence_level=0.95
    )
    
    # Peak demand prediction
    peak_demand = consumption_forecast.predictions.max()
    peak_hour = consumption_forecast.predictions.argmax()
    
    # Load balancing recommendations
    load_balancing = {
        'peak_demand': peak_demand,
        'peak_hour': peak_hour,
        'recommended_generation': peak_demand * 1.1,  # 10% buffer
        'demand_response_potential': peak_demand * 0.15  # 15% reduction potential
    }
    
    return consumption_forecast, load_balancing

# Renewable energy forecasting
def forecast_renewable_generation(weather_data, generation_data, horizon_hours=24):
    ts_model = TabPFNTimeSeries()
    
    # Multi-variate forecasting with weather data
    renewable_forecast = ts_model.forecast_multivariate(
        data=pd.concat([weather_data, generation_data], axis=1),
        horizon=horizon_hours,
        series_names=['solar_irradiance', 'wind_speed', 'solar_generation', 'wind_generation'],
        confidence_level=0.90
    )
    
    return renewable_forecast

# Patient flow forecasting
def forecast_patient_flow(admission_data, horizon_days=7):
    ts_model = TabPFNTimeSeries()
    
    # Weekly patient flow forecast
    flow_forecast = ts_model.forecast(
        data=admission_data,
        horizon=horizon_days,
        confidence_level=0.90
    )
    
    # Resource planning
    resource_planning = {
        'predicted_admissions': flow_forecast.predictions.sum(),
        'bed_requirements': flow_forecast.predictions.sum() * 1.2,  # 20% buffer
        'staffing_needs': flow_forecast.predictions.sum() * 0.1,  # 1 staff per 10 patients
        'confidence_level': 0.90
    }
    
    return flow_forecast, resource_planning

# Disease outbreak forecasting
def forecast_disease_spread(incidence_data, horizon_days=30):
    ts_model = TabPFNTimeSeries()
    
    # Disease spread forecast
    spread_forecast = ts_model.forecast(
        data=incidence_data,
        horizon=horizon_days,
        confidence_level=0.95
    )
    
    # Public health recommendations
    public_health_planning = {
        'predicted_cases': spread_forecast.predictions.sum(),
        'peak_incidence': spread_forecast.predictions.max(),
        'peak_date': spread_forecast.predictions.argmax(),
        'intervention_threshold': spread_forecast.predictions.mean() * 2
    }
    
    return spread_forecast, public_health_planning