# tools/math_tools.py

import numpy as np
import pandas as pd
from scipy.optimize import linprog
from statsmodels.tsa.arima.model import ARIMA
import matplotlib.pyplot as plt

def run_math_analysis(topic: str):
    if topic == "linear_algebra":
        linear_algebra_demo()
    elif topic == "probability":
        probability_demo()
    elif topic == "optimization":
        optimization_demo()
    elif topic == "game_theory":
        game_theory_demo()
    elif topic == "time_series":
        time_series_demo()
    elif topic == "markov":
        markov_process_demo()
    else:
        print("❌ Unknown topic. Available: linear_algebra, probability, optimization, game_theory, time_series, markov")

def linear_algebra_demo():
    print("📐 Linear Algebra: Matrix-Vector Multiplication")
    A = np.array([[1, 2], [3, 4]])
    B = np.array([[2], [1]])
    result = np.dot(A, B)
    print("Matrix A:\n", A)
    print("Vector B:\n", B)
    print("A · B:\n", result)

def probability_demo():
    print("📊 Probability & Bayesian Inference Example")
    prior = 0.3
    likelihood = 0.8
    evidence = 0.6
    posterior = (likelihood * prior) / evidence
    print(f"Bayesian Posterior = P(H|E) = {posterior:.2f}")

def optimization_demo():
    print("🔧 Optimization: Linear Programming (Maximize 3x + 5y)")
    c = [-3, -5]  # Maximize 3x + 5y → minimize -3x -5y
    A = [[1, 0], [0, 2], [3, 2]]
    b = [4, 12, 18]
    result = linprog(c, A_ub=A, b_ub=b)
    if result.success:
        print("Optimal values of x and y:", result.x)
        print("Maximum value of objective function:", -result.fun)
    else:
        print("❌ Optimization failed:", result.message)

def game_theory_demo():
    print("🎲 Game Theory: Payoff Matrix")
    matrix = np.array([[3, 0], [5, 1]])
    print("Player A vs Player B Payoff Matrix:")
    print(matrix)
    print("Interpret rows as Player A strategies and columns as Player B.")

def time_series_demo():
    print("📈 Time Series: ARIMA(1,1,1) Forecast")
    np.random.seed(42)
    series = pd.Series(np.random.randn(100).cumsum())
    model = ARIMA(series, order=(1, 1, 1))
    fit = model.fit()
    forecast = fit.forecast(steps=10)
    print("Forecasted next 10 values:")
    print(forecast)

    series.plot(label="Historical")
    forecast.plot(label="Forecast", legend=True)
    plt.title("Time Series Forecast (ARIMA)")
    plt.xlabel("Time")
    plt.ylabel("Value")
    plt.grid(True)
    plt.show()

def markov_process_demo():
    print("🔁 Markov Chain Simulation: Sunny vs Rainy")
    transition_matrix = np.array([[0.9, 0.1], [0.5, 0.5]])  # [Sunny, Rainy]
    state = np.array([1, 0])  # Start in Sunny
    for day in range(5):
        state = np.dot(state, transition_matrix)
        print(f"Day {day + 1}: P(Sunny) = {state[0]:.2f}, P(Rainy) = {state[1]:.2f}")