import matplotlib.pyplot as plt
import numpy as np
from numpy.polynomial import Polynomial
import sympy as sp
import json

def calc_payment_MSE(polynomial, data, divisor):
    data_np = np.array(data)  # shape (N,2)
    xs = data_np[:, 0]
    ys = data_np[:, 1]

    # Find unique x values and their indices
    unique_xs, inverse_indices = np.unique(xs, return_inverse=True)

    # Evaluate polynomial once for each unique x
    poly_vals_unique = polynomial(unique_xs)
    poly_vals_unique = np.clip(poly_vals_unique, 0, 10)

    # Map back to full data length
    poly_vals = poly_vals_unique[inverse_indices]

    # Compute squared error
    sq_errors = (poly_vals - ys) ** 2

    # Calculate payments vectorized
    # payments = np.where(sq_errors < 100, 100 - sq_errors, 0)
    # payments = payments/100*25
    payments = 25*(1-sq_errors/divisor)

    payment = payments.mean()
    payment = max(payment, 0)

    print(f"Payment: {payment:.2f}")
    return round(payment, 2)

def heatmaps(polynomial,divisor):
    x = np.linspace(0, 10, 50)
    y = np.linspace(0, 10, 50)
    data = []
    for i in range(len(x)):
        for j in range(len(y)):
            poly_value = polynomial(x[i])
            poly_value = np.clip(poly_value, 0, 10)
            payment = 25*(1-(poly_value - y[j])**2/divisor)
            data.append([[x[i], y[j]], round(payment, 2)])
    return data

def calc_payment_MSE_data(points, data, divisor):
    if isinstance(points, str):
        points = json.loads(points)
    if isinstance(data, str):
        data = json.loads(data)

    data_np = np.array(data)   # shape (N, 2)

    points_np = np.array(points)  # shape (M, 2)

    xs_data = data_np[:, 0]
    ys_data = data_np[:, 1]
    xs_points = points_np[:, 0]
    ys_points = points_np[:, 1]

    # For each x in data, find the index of the closest x in points
    indices = np.abs(xs_points[:, None] - xs_data[None, :]).argmin(axis=0)

    # Retrieve the corresponding y values
    matched_ys = ys_points[indices]

    # Compute squared errors
    sq_errors = (matched_ys - ys_data) ** 2

    # Compute MSE
    # payments = np.where(sq_errors < 100, 100 - sq_errors, 0)
    # payments = payments/100*25
    payments = 25*(1-sq_errors/divisor)

    payment = payments.mean()
    payment = max(payment, 0)

    print(f"Payment: {payment:.2f}")
    return round(payment, 2)

def heatmaps_data(coordinates, divisor):
    x = np.linspace(0, 10, 50)
    y = np.linspace(0, 10, 50)
    data = []
    for i in range(len(x)):
        for j in range(len(y)):
            points_np = np.array(coordinates)
            xs_points = points_np[:, 0]
            ys_points = points_np[:, 1]
            index = np.abs(xs_points - x[i]).argmin()
            anno_value = ys_points[index]
            anno_value = np.clip(anno_value, 0, 10)
            payment = 25*(1-(anno_value - y[j])**2/divisor)
            data.append([[x[i], y[j]], round(payment, 2)])
    return data