from fastmcp import FastMCP
import numpy as np
from pydantic import BaseModel
from typing import List, Tuple, Optional
from scipy import stats

mcp = FastMCP("Demo 🚀")

class HelloInput(BaseModel):
    name: str

@mcp.tool()
def hello(input: HelloInput) -> str:
    return f"Hello, {input.name}!"

class MultiplyInput(BaseModel):
    a: float
    b: float

@mcp.tool()
def multiply(input: MultiplyInput) -> float:
    """Multiplies two numbers."""
    return input.a * input.b

class InnerProductInput(BaseModel):
    a: List[float]
    b: List[float]

@mcp.tool()
def inner_product(input: InnerProductInput) -> float:
    """Calculates the inner product of two vectors."""
    return np.dot(input.a, input.b)

class MatrixMultiplyInput(BaseModel):
    a: List[List[float]]
    b: List[List[float]]

@mcp.tool()
def matrix_multiply(input: MatrixMultiplyInput) -> List[List[float]]:
    """Multiplies two matrices."""
    return np.matmul(input.a, input.b)

class NumpyDotInput(BaseModel):
    a: List[float]
    b: List[float]

@mcp.tool()
def numpy_dot(input: NumpyDotInput) -> float:
    """Calculates the dot product of two vectors."""
    return np.dot(input.a, input.b)

class NumpyMatmulInput(BaseModel):
    a: List[List[float]]
    b: List[List[float]]

@mcp.tool()
def numpy_matmul(input: NumpyMatmulInput) -> List[List[float]]:
    """Multiplies two matrices using matmul."""
    return np.matmul(input.a, input.b)

class NumpyInvInput(BaseModel):
    a: List[List[float]]

@mcp.tool()
def numpy_inv(input: NumpyInvInput) -> List[List[float]]:
    """Calculates the inverse of a matrix."""
    return np.linalg.inv(input.a)

class NumpyDetInput(BaseModel):
    a: List[List[float]]

@mcp.tool()
def numpy_det(input: NumpyDetInput) -> float:
    """Calculates the determinant of a matrix."""
    return np.linalg.det(input.a)

class NumpyEigInput(BaseModel):
    a: List[List[float]]

@mcp.tool()
def numpy_eig(input: NumpyEigInput) -> Tuple:
    """Calculates the eigenvalues and eigenvectors of a matrix."""
    return np.linalg.eig(input.a)

class NumpySvdInput(BaseModel):
    a: List[List[float]]

@mcp.tool()
def numpy_svd(input: NumpySvdInput) -> Tuple:
    """Performs singular value decomposition on a matrix."""
    return np.linalg.svd(input.a)

class NumpyNormInput(BaseModel):
    a: List[float]
    ord: Optional[int] = None

@mcp.tool()
def numpy_norm(input: NumpyNormInput) -> float:
    """Calculates the norm of a vector or matrix."""
    return np.linalg.norm(input.a, input.ord)

class NumpyCrossInput(BaseModel):
    a: List[float]
    b: List[float]

@mcp.tool()
def numpy_cross(input: NumpyCrossInput) -> List[float]:
    """Calculates the cross product of two vectors."""
    return np.cross(input.a, input.b)

class NumpyInnerInput(BaseModel):
    a: List[float]
    b: List[float]

@mcp.tool()
def numpy_inner(input: NumpyInnerInput) -> float:
    """Calculates the inner product of two vectors."""
    return np.inner(input.a, input.b)

class NumpyOuterInput(BaseModel):
    a: List[float]
    b: List[float]

@mcp.tool()
def numpy_outer(input: NumpyOuterInput) -> List[List[float]]:
    """Calculates the outer product of two vectors."""
    return np.outer(input.a, input.b)

class NumpyTensordotInput(BaseModel):
    a: List
    b: List
    axes: int = 2

@mcp.tool()
def numpy_tensordot(input: NumpyTensordotInput) -> float:
    """Calculates the tensor dot product of two arrays."""
    return np.tensordot(input.a, input.b, input.axes)

class NumpyTraceInput(BaseModel):
    a: List[List[float]]

@mcp.tool()
def numpy_trace(input: NumpyTraceInput) -> float:
    """Calculates the trace of a matrix."""
    return np.trace(input.a)

class NumpyQrInput(BaseModel):
    a: List[List[float]]

@mcp.tool()
def numpy_qr(input: NumpyQrInput) -> Tuple:
    """Performs QR decomposition on a matrix."""
    return np.linalg.qr(input.a)

class NumpyCholeskyInput(BaseModel):
    a: List[List[float]]

@mcp.tool()
def numpy_cholesky(input: NumpyCholeskyInput) -> List[List[float]]:
    """Performs Cholesky decomposition on a matrix."""
    return np.linalg.cholesky(input.a)

class NumpySolveInput(BaseModel):
    a: List[List[float]]
    b: List[float]

@mcp.tool()
def numpy_solve(input: NumpySolveInput) -> List[float]:
    """Solves a linear matrix equation."""
    return np.linalg.solve(input.a, input.b)

class NumpyLstsqInput(BaseModel):
    a: List[List[float]]
    b: List[float]

@mcp.tool()
def numpy_lstsq(input: NumpyLstsqInput) -> Tuple:
    """Solves a linear least squares problem."""
    return np.linalg.lstsq(input.a, input.b, rcond=None)

class NumpyPinvInput(BaseModel):
    a: List[List[float]]

@mcp.tool()
def numpy_pinv(input: NumpyPinvInput) -> List[List[float]]:
    """Calculates the Moore-Penrose pseudo-inverse of a matrix."""
    return np.linalg.pinv(input.a)

class NumpyCondInput(BaseModel):
    a: List[List[float]]
    p: Optional[int] = None

@mcp.tool()
def numpy_cond(input: NumpyCondInput) -> float:
    """Calculates the condition number of a matrix."""
    return np.linalg.cond(input.a, input.p)

class NumpyMatrixRankInput(BaseModel):
    a: List[List[float]]

@mcp.tool()
def numpy_matrix_rank(input: NumpyMatrixRankInput) -> int:
    """Calculates the rank of a matrix."""
    return np.linalg.matrix_rank(input.a)

class NumpyMultiDotInput(BaseModel):
    arrays: List[List[List[float]]]

@mcp.tool()
def numpy_multi_dot(input: NumpyMultiDotInput) -> List[List[float]]:
    """Computes the dot product of two or more arrays in a single function call, while automatically selecting the fastest evaluation order."""
    return np.linalg.multi_dot(input.arrays)

# Static resource
@mcp.resource("config://version")
def get_version(): 
    return "2.0.1"

@mcp.resource("users://{user_id}/profile")
def get_profile(user_id: int):
    # Fetch profile for user_id...
    return {"name": f"User {user_id}", "status": "active"}

class SummarizeRequestInput(BaseModel):
    text: str

@mcp.prompt()
def summarize_request(input: SummarizeRequestInput) -> str:
    """Generate a prompt asking for a summary."""
    return f"Please summarize the following text:\n\n{input.text}"

class ScipyTtestInput(BaseModel):
    a: List[float]
    b: List[float]

@mcp.tool()
def scipy_ttest(input: ScipyTtestInput) -> dict:
    """Performs an independent two-sample t-test."""
    t_stat, p_value = stats.ttest_ind(input.a, input.b)
    return {"t_statistic": t_stat, "p_value": p_value}

class ScipyPearsonrInput(BaseModel):
    x: List[float]
    y: List[float]

@mcp.tool()
def scipy_pearsonr(input: ScipyPearsonrInput) -> dict:
    """Calculates the Pearson correlation coefficient."""
    corr_coefficient, p_value = stats.pearsonr(input.x, input.y)
    return {"correlation_coefficient": corr_coefficient, "p_value": p_value}

if __name__ == "__main__":
    mcp.run(transport="sse", host="0.0.0.0", port=7860)