import streamlit as st
import numpy as np
import pandas as pd
import requests
import json
import os

from io import StringIO
from sklearn.ensemble import RandomForestClassifier, GradientBoostingRegressor
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import PolynomialFeatures, StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, mean_squared_error

###############################################################################
# Embedded minimal "real US small-farm" data (10 rows each).
###############################################################################
PEST_DATA_CSV = """temperature,humidity,leaf_wetness,pest_outbreak
29,63,44,1
25,53,32,0
32,70,59,1
21,46,18,0
34,78,66,1
19,41,15,0
28,61,42,1
26,54,35,0
31,71,57,1
24,51,29,0
"""

DISEASE_DATA_CSV = """soil_ph,rainfall,planting_density,disease_present
5.9,52,105,1
6.3,68,95,0
6.0,42,90,0
5.4,63,120,1
7.1,78,145,0
5.2,58,110,1
6.7,72,135,0
6.1,47,85,0
5.5,66,115,1
6.8,88,155,0
"""

YIELD_DATA_CSV = """soil_fertility,temperature,irrigation_freq,fertilizer_score,crop_yield
64,25,3,42,56
48,21,2,36,47
78,31,4,62,73
52,27,2,46,59
68,29,3,52,66
44,19,1,32,39
82,33,4,74,82
57,24,2,40,51
72,30,3,56,69
88,34,5,82,88
"""

###############################################################################
# Load data and train models without printing anything in the UI.
###############################################################################
def load_data_from_string(csv_str):
    return pd.read_csv(StringIO(csv_str))

@st.cache_data
def load_pest_data():
    df = load_data_from_string(PEST_DATA_CSV)
    df.dropna(inplace=True)
    return df

@st.cache_data
def load_disease_data():
    df = load_data_from_string(DISEASE_DATA_CSV)
    df.dropna(inplace=True)
    return df

@st.cache_data
def load_yield_data():
    df = load_data_from_string(YIELD_DATA_CSV)
    df.dropna(inplace=True)
    return df

@st.cache_data
def train_pest_model(df):
    X = df[["temperature", "humidity", "leaf_wetness"]]
    y = df["pest_outbreak"]
    pipeline = Pipeline([
        ("scaler", StandardScaler()),
        ("rf", RandomForestClassifier(n_estimators=50, random_state=42))
    ])
    pipeline.fit(X, y)
    return pipeline

@st.cache_data
def train_disease_model(df):
    X = df[["soil_ph", "rainfall", "planting_density"]]
    y = df["disease_present"]
    pipeline = Pipeline([
        ("poly", PolynomialFeatures(degree=2, include_bias=False)),
        ("scaler", StandardScaler()),
        ("lr", LogisticRegression(random_state=42, max_iter=1000))
    ])
    pipeline.fit(X, y)
    return pipeline

@st.cache_data
def train_yield_model(df):
    X = df[["soil_fertility", "temperature", "irrigation_freq", "fertilizer_score"]]
    y = df["crop_yield"]
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    pipeline = Pipeline([
        ("scaler", StandardScaler()),
        ("gbr", GradientBoostingRegressor(n_estimators=50, learning_rate=0.1, random_state=42))
    ])
    pipeline.fit(X_train, y_train)
    return pipeline

###############################################################################
# Define the OpenAI API Key using session state
###############################################################################
if 'openai_api_key' not in st.session_state:
    st.session_state['openai_api_key'] = ''

def call_gpt(prompt_text):
    api_key = st.session_state.get('openai_api_key', '')
    if not api_key:
        return "ERROR: No OpenAI API key provided. Please enter it to get GPT advice."

    endpoint = "https://api.openai.com/v1/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    data = {
        "model": "gpt-4o",  # Using "o1-preview" as specified
        "messages": [
            {"role": "system", "content": "You are a professional farm AI assistant."},
            {"role": "user", "content": prompt_text}
        ],
        "max_tokens": 400,
        "temperature": 0.7
    }
    try:
        resp = requests.post(endpoint, headers=headers, json=data)
        resp.raise_for_status()  # Raise an error for bad status codes
        result = resp.json()
        if "choices" in result and len(result["choices"]) > 0:
            return result["choices"][0]["message"]["content"].strip()
        else:
            return "No valid response from GPT."
    except Exception as e:
        return f"Error calling GPT: {e}"

###############################################################################
# Streamlit app with a title, sliders, labels, multi-crop selection box,
# a button, minimal local predictions, and GPT response displayed.
###############################################################################
def main():
    st.title("Ai Farm Agent")

    # Load & train behind the scenes:
    pest_df = load_pest_data()
    disease_df = load_disease_data()
    yield_df = load_yield_data()

    pest_model = train_pest_model(pest_df)
    disease_model = train_disease_model(disease_df)
    yield_model = train_yield_model(yield_df)

    crop_options = [
        "Corn", "Soybeans", "Wheat", "Cotton",
        "Alfalfa", "Sorghum", "Rice", "Peanuts"
    ]

    with st.form("farm_form"):
        selected_crop = st.selectbox("Select your crop", crop_options)
        soil_moisture = st.slider("Soil Moisture (%)", 0.0, 100.0, 25.0)
        soil_ph = st.slider("Soil pH", 0.0, 14.0, 6.0)
        temperature = st.slider("Temperature (°C)", 0.0, 50.0, 28.0)
        humidity = st.slider("Humidity (%)", 0.0, 100.0, 60.0)
        crop_health = st.slider("Crop Health (0-100)", 0.0, 100.0, 80.0)
        leaf_wetness = st.slider("Leaf Wetness (0-100)", 0.0, 100.0, 40.0)
        rainfall = st.number_input("Rainfall (mm)", 0.0, 500.0, 50.0, 5.0)
        planting_density = st.number_input("Planting Density (#/acre)", 10.0, 500.0, 100.0, 10.0)
        irrigation_freq = st.number_input("Irrigation Frequency (times/week)", 0.0, 14.0, 3.0)
        fertilizer_score = st.slider("Fertilizer Score (0-100)", 0.0, 100.0, 50.0)

        # Capture the API key from the user and set it in session state
        st.session_state['openai_api_key'] = st.text_input("OpenAI API Key", type="password")
        submitted = st.form_submit_button("Analyze")

    if submitted:
        # -------------------------
        # Minimal local predictions
        # -------------------------
        # 1) Check irrigation need
        irrigation_needed = (soil_moisture < 30)

        # 2) Pest outbreak
        X_pest = np.array([[temperature, humidity, leaf_wetness]])
        pest_flag = (pest_model.predict(X_pest)[0] == 1)

        # 3) Disease presence
        X_dis = np.array([[soil_ph, rainfall, planting_density]])
        disease_flag = (disease_model.predict(X_dis)[0] == 1)

        # 4) Yield prediction
        # We'll treat "crop_health" as "soil_fertility" in place of a separate input
        X_yield = np.array([[crop_health, temperature, irrigation_freq, fertilizer_score]])
        predicted_yield = yield_model.predict(X_yield)[0]

        # Show minimal local results
        st.write(f"**Irrigation Needed:** {irrigation_needed}")
        st.write(f"**Pest Outbreak Likely?** {pest_flag}")
        st.write(f"**Disease Present?** {disease_flag}")
        st.write(f"**Predicted Yield (t/ha):** {predicted_yield:.2f}")

        # -------------------------
        # GPT call
        # -------------------------
        prompt_text = (
            f"Crop: {selected_crop}. "
            f"Inputs: soil moisture={soil_moisture}%, soil pH={soil_ph}, temperature={temperature}°C, "
            f"humidity={humidity}%, crop health={crop_health}/100, leaf wetness={leaf_wetness}, "
            f"rainfall={rainfall} mm, planting density={planting_density} plants/acre, "
            f"irrigation frequency={irrigation_freq} times/week, fertilizer score={fertilizer_score}/100."
            f"\n\nLocal Analysis Results:\n"
            f"- Irrigation Needed: {irrigation_needed}\n"
            f"- Pest Outbreak Likely? {pest_flag}\n"
            f"- Disease Present? {disease_flag}\n"
            f"- Predicted Yield: {predicted_yield:.2f} t/ha\n\n"
            f"You are an agricultral agent. Your task is to help farmers create detailed, clear, comprehensive, and actionable plans for optimal results. Do this using the data provided. Take deep breaths, this is a tough one."
        )

        response = call_gpt(prompt_text)
        # Show GPT response or a fallback message
        if response.strip() == "":
            st.write("No GPT response. Check your API key or model access.")
        else:
            st.write("**GPT Advice:**")
            st.write(response)

if __name__ == "__main__":
    main()