from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline

app = FastAPI()

class CalculationRequest(BaseModel):
    a: float
    b: float
    operation: str

model = pipeline('text-generation', model='distilgpt2')

@app.get("/")
def read_root():
    return {"message": "Welcome to the Math API"}

@app.post("/calculate")
def calculate(request: CalculationRequest):
    a = request.a
    b = request.b
    operation = request.operation.lower()

    try:
        if operation == "add":
            result = a + b
        elif operation == "subtract":
            result = a - b
        elif operation == "multiply":
            result = a * b
        elif operation == "divide":
            if b == 0:
                raise HTTPException(status_code=400, detail="Division by zero is not allowed")
            result = a / b
        elif operation == "power":
            result = math.pow(a, b)
        elif operation == "modulus":
            result = a % b
        elif operation == "sqrt":
            if a < 0:
                raise HTTPException(status_code=400, detail="Square root of negative number is not allowed")
            result = math.sqrt(a)
        elif operation == "log":
            if a <= 0:
                raise HTTPException(status_code=400, detail="Logarithm of non-positive number is not allowed")
            result = math.log(a)
        elif operation == "log10":
            if a <= 0:
                raise HTTPException(status_code=400, detail="Logarithm of non-positive number is not allowed")
            result = math.log10(a)
        elif operation == "sin":
            result = math.sin(math.radians(a))
        elif operation == "cos":
            result = math.cos(math.radians(a))
        elif operation == "tan":
            result = math.tan(math.radians(a))
        elif operation == "asin":
            result = math.degrees(math.asin(a))
        elif operation == "acos":
            result = math.degrees(math.acos(a))
        elif operation == "atan":
            result = math.degrees(math.atan(a))
        elif operation == "factorial":
            if a < 0 or not float(a).is_integer():
                raise HTTPException(status_code=400, detail="Factorial of negative or non-integer number is not allowed")
            result = math.factorial(int(a))
        else:
            raise HTTPException(status_code=400, detail="Invalid operation")
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

    return {"result": result}

@app.post("/generate")
def generate_text(prompt: str):
    generated = model(prompt, max_length=50, clean_up_tokenization_spaces=True)
    return {"generated_text": generated[0]['generated_text']}

@app.post("/test_math")
def test_math(request: CalculationRequest):
    a = request.a
    b = request.b
    operation = request.operation

    if operation == "add":
        result = a + b
    elif operation == "subtract":
        result = a - b
    elif operation == "multiply":
        result = a * b
    elif operation == "divide":
        result = a / b
    else:
        return {"error": "Invalid operation"}

    return {"result": result}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)