import requests
import pandas as pd
import numpy as np


def main(loc_id):
    """
    Fetch survey history for a fish farm location ID and expand with measurement data
    for report_type 'b' (b-undersøkelser).

    Adds columns:
        - lat (list)
        - lon (list)
        - b_float_value (list)
        - tilstand_final (list)
        - depth (list)
        - measurement number (list)

    For reports without measurement data (including type 'c'),
    these columns contain empty lists.
    """

    loc_id = str(loc_id)

    base_url = "https://api.fiskeridir.no/envreportreg-public/api/v2/report/search"
    params = {"siteNumber": loc_id}

    response = requests.get(base_url, params=params)
    response.raise_for_status()

    data = response.json()
    items = data.get("content", [])

    if not items:
        return pd.DataFrame(columns=[
            "report_type", "report_date", "report_id", "download_link",
            "lat", "lon", "b_float_value", "tilstand_final"
        ])
    dl_link = "https://api.fiskeridir.no/envreportreg-public/api/v2/report/"

    # Build base dataframe
    df = pd.DataFrame({
        "report_type": [item.get("envExaminationType") for item in items],
        "report_date": [item.get("reportCreated") for item in items],
        "report_id":   [int(item.get("reportId")) for item in items],
        "download_link" : [dl_link + str(item.get("reportId")) +  '/pdf' for item in items],
    })

    # Convert and sort
    df["report_date"] = pd.to_datetime(df["report_date"], utc=True)

    df = df.sort_values(
        by=["report_type", "report_date"],
        ascending=[True, False]
    ).reset_index(drop=True)

    # Prepare empty columns (lists per row)
    df["lat"] = [[] for _ in range(len(df))]
    df["lon"] = [[] for _ in range(len(df))]
    df["b_float_value"] = [[] for _ in range(len(df))]
    df["tilstand_final"] = [[] for _ in range(len(df))]
    df["depth"] = [[] for _ in range(len(df))]
    df["measurement_nr"] = [[] for _ in range(len(df))]

    # Process only report_type 'b'
    for idx, row in df.iterrows():

        if row["report_type"] != "b":
            continue

        report_id = str(row["report_id"])
        detail_url = f"https://api.fiskeridir.no/envreportreg-public/api/v2/report/{report_id}"

        response = requests.get(detail_url)
        response.raise_for_status()

        detail_data = response.json()
        measurements = detail_data.get("measurements", [])

        if not measurements:
            continue  # leave empty lists

        # Extract values
        lat = np.array([m.get("latitudeDecimal") for m in measurements])
        lon = np.array([m.get("longitudeDecimal") for m in measurements])
        b_float_value = np.array([m.get("medianGroupIIAndIIIValue") for m in measurements])
        depth = np.array([m.get("depthMeters") for m in measurements])
        measurement_nr = np.array([m.get("stationMeasurementNumber") for m in measurements])

        # Keep only valid coordinates
        valid_idx = np.where(lat > 0)[0]

        lat = lat[valid_idx]
        lon = lon[valid_idx]
        b_float_value = b_float_value[valid_idx]
        depth = depth[valid_idx]
        measurement_nr = measurement_nr[valid_idx]

        # Compute tilstand_final (vectorized)
        tilstand_final = np.select(
            [
                b_float_value < 1.1,
                (b_float_value >= 1.1) & (b_float_value < 2.1),
                (b_float_value >= 2.1) & (b_float_value < 3.1),
            ],
            [1, 2, 3],
            default=4
        )

        # Store as lists in dataframe
        df.at[idx, "lat"] = lat.tolist()
        df.at[idx, "lon"] = lon.tolist()
        df.at[idx, "b_float_value"] = b_float_value.tolist()
        df.at[idx, "tilstand_final"] = tilstand_final.tolist()
        df.at[idx, "depth"] = depth.tolist()
        df.at[idx, "measurement_nr"] = measurement_nr.tolist()

    # Convert date back to string if desired
    df["report_date"] = df["report_date"].dt.strftime("%Y-%m-%dT%H:%M:%SZ")


    '''Example for plotting b_survey results:
import matplotlib.pyplot as plt
survey = 2
color_map = {
    1: 'blue',
    2: 'green',
    3: 'yellow',
    4: 'red'
}

#colors = pd['tilstand_final'][0].map(color_map)
colors = pd.Series(df2['tilstand_final'][survey]).map(color_map).tolist()

plt.figure()

plt.scatter(
    df2['lon'][survey],
    df2['lat'][survey],
    c=colors,
    marker='s'  # square
)

plt.xlabel('Longitude')
plt.ylabel('Latitude')
plt.title('Tilstand Plot B rapport '+ df2['report_date'][survey])

plt.show()'''

    return df