#File system and I/O operations
import os
from io import BytesIO
# Data manipulation 
import pandas as pd
import numpy as np
from IPython.display import Image

# Time functions
import math
import time
from datetime import datetime, timedelta

# Data Gathering
import requests
import zipfile

#Encoding and Serialization
import base64
import json

#Debugging
import http
http.client.HTTPConnection.debuglevel = 0

#Visualization Libraries
import plotly.graph_objects as go
import plotly.express as px

# Mention your working directory
os.chdir("C:\\Users\\XX")

## Dataset Libraries
from pyobis import dataset
from pyobis import occurrences

#pelagis_df.columns.to_list()
#We pick the ones that are relevant to our use case
columns_considered = [
 "eventDate",
 "occurrenceID",
 "scientificName",
 "genus",
 "species",
 "aphiaID",
 "family",
 "class",
 "vernacularName",
 "dateIdentified",
 "bathymetry",
 "shoredistance",
 "sst",
 "sss",
 "marine",
 "brackish",
 "absence",
 "flags",
 "waterBody",
 "occurrenceStatus"  
]
pelagis_df_final = pelagis_df[columns_considered]

#Standardizing Column Datatypes
pelagis_df_final['flags'] = pelagis_df_final['flags'].astype(str)
pelagis_df_final['brackish'] = pelagis_df_final['brackish'].astype(str)

username = "XX" # Provide your GridDB credentials here
password =   "XX"
credentials = f"{username}:{password}"
encoded_credentials = base64.b64encode(credentials.encode()).decode()
#print(f"Encoded credentials: Basic {encoded_credentials}")

#Construct an object to hold the request headers (ensure that you replace the XXX placeholder with the correct value that matches the credentials for your GridDB instance)
#header_obj = {"Authorization":"Basic {EncodedAuthorizationToken}","Content-Type":"application/json; charset=UTF-8","User-Agent":"PostmanRuntime/7.29.0"}

header_obj = {
    "Authorization": f"Basic {encoded_credentials}",  # Add encoded credentials here
    "Content-Type": "application/json",  # Optional; depends on API requirements
    "charset": "UTF-8",
    "User-Agent":"PostmanRuntime/7.29.0"
}

#Construct the base URL based on your GridDB cluster you'd like to connect to (ensure that you replace the placeholders in the URL below with the correct values that correspond to your GridDB instance)
base_url = https://[host]:[port]/griddb/v2/[clustername]/dbs/[database_name]

def map_dtype_to_griddb(dtype):
 # Maps Pandas data types to GridDB data types.
    if pd.api.types.is_integer_dtype(dtype):
        return "INTEGER"
    elif pd.api.types.is_float_dtype(dtype):
        return "DOUBLE"
    elif pd.api.types.is_bool_dtype(dtype):
        return "BOOL"
    ##elif pd.api.types.is_datetime64_any_dtype(dtype):
        ##return "LONG"  # GridDB stores timestamps as LONG
    elif pd.api.types.is_datetime64_any_dtype(dtype):
        return "TIMESTAMP"
    else:
        return "STRING"

def generate_griddb_data_obj(df, container_name="MyContainer", container_type="COLLECTION", rowkey=False):
    #Generates a GridDB container data object for API request.
    columns = []
    for col in df.columns:
        griddb_type = map_dtype_to_griddb(df[col].dtype)
        columns.append({"name": col, "type": griddb_type})
    
    data_obj = {
        "container_name": container_name,
        "container_type": container_type,
        "rowkey": rowkey,
        "columns": columns
    }
    
    return data_obj

#Setup the URL to be used to invoke the GridDB WebAPI to register rows in the container created previously
url = base_url + 'containers/pelagis/rows'
print(url)

batch_size = 5000
total_rows = len(pelagis_df_final)
total_batches = math.ceil(total_rows / batch_size)

print(f"Starting ingestion of {total_rows:,} rows in {total_batches} batches of {batch_size} each...")
start_time = time.time()

for i in range(total_batches):
    start_idx = i * batch_size
    end_idx = min((i + 1) * batch_size, total_rows)
    batch_df = pelagis_df_final.iloc[start_idx:end_idx]
    batch_json = batch_df.to_json(orient='values')
    
    try:
        print(url)
        response = requests.put(url, headers=header_obj, data=batch_json)
        elapsed = time.time() - start_time

        if response.status_code in [200, 201]:
            print(f" Batch {i+1}/{total_batches} | Rows {start_idx}-{end_idx-1} | "
                  f"Time elapsed: {elapsed:.1f}s")
        else:
            print(f" Batch {i+1}/{total_batches} FAILED "
                  f"(HTTP {response.status_code}) - {response.text[:150]}")

    except Exception as e:
        print(f" Batch {i+1}/{total_batches} ERROR: {e}")

print(f"\n Ingestion completed in {time.time() - start_time:.1f} seconds.")

sql_query1 = ("""
            SELECT p.genus, 
            p.class,     
            EXTRACT(YEAR, TIMESTAMP(p.dateIdentified)) AS year_identified,
            p.vernacularName, 
            count(DISTINCT p.occurrenceID) as num_occurrences
            FROM pelagis p
            GROUP BY 1,2,3,4
            HAVING COUNT(DISTINCT p.occurrenceID) > 0
            ORDER BY 5 DESC
            """)

#Setup the URL to be used to invoke the GridDB WebAPI to retrieve data from the container
url = base_url + 'sql'

# Construct the request body
request_body = json.dumps([
    {
        "type": "sql-select",
        "stmt": sql_query1
    }
])

# Disable HTTP debugging
http.client.HTTPConnection.debuglevel = 0

#Invoke the GridDB WebAPI
data_req1 = requests.post(url, data=request_body, headers=header_obj)

#Process the response received and construct a Pandas dataframe with the data from the response
myJson = data_req1.json()
occurrences_df = pd.DataFrame(myJson[0]["results"], columns=[myJson[0]["columns"][0]["name"], myJson[0]["columns"][1]["name"], myJson[0]["columns"][2]["name"], myJson[0]["columns"][3]["name"], myJson[0]["columns"][4]["name"]])

sql_query2 = ("""
                SELECT
                EXTRACT(YEAR, TIMESTAMP(p.eventDate)) AS year,
                AVG(p.sst) AS avg_sst,
                AVG(p.sss) AS avg_sss
                FROM pelagis p
                WHERE p.occurrenceStatus = 'present'
                GROUP BY year
                ORDER BY year
                """)

#Setup the URL to be used to invoke the GridDB WebAPI to retrieve data from the container
url = base_url + 'sql'

# Construct the request body
request_body = json.dumps([
    {
        "type": "sql-select",
        "stmt": sql_query2
    }
])

# Disable HTTP debugging
http.client.HTTPConnection.debuglevel = 0

#Invoke the GridDB WebAPI
data_req2 = requests.post(url, data=request_body, headers=header_obj)

#Process the response received and construct a Pandas dataframe with the data from the response
myJson = data_req2.json()
avg_metrics_df = pd.DataFrame(myJson[0]["results"], columns=[myJson[0]["columns"][0]["name"], myJson[0]["columns"][1]["name"], myJson[0]["columns"][2]["name"]])

sql_query3 = ("""SELECT
            CASE 
                WHEN p.bathymetry < 500 THEN 'Shallow (<500m)'
                WHEN p.bathymetry BETWEEN 500 AND 1500 THEN 'Mid-depth (500-1500m)'
                ELSE 'Deep (>1500m)'
            END AS depth_range,
            COUNT(p.occurrenceID) AS num_observations
            FROM pelagis p
            WHERE p.occurrenceStatus = 'present'
            GROUP BY depth_range
            ORDER BY 2 DESC """)

#Setup the URL to be used to invoke the GridDB WebAPI to retrieve data from the container
url = base_url + 'sql'

# Construct the request body
request_body = json.dumps([
    {
        "type": "sql-select",
        "stmt": sql_query3
    }
])

# Disable HTTP debugging
http.client.HTTPConnection.debuglevel = 0

#Invoke the GridDB WebAPI
data_req3 = requests.post(url, data=request_body, headers=header_obj)

#Process the response received and construct a Pandas dataframe with the data from the response
myJson = data_req3.json()
occurrences_by_depth_range_df = pd.DataFrame(myJson[0]["results"], columns=[myJson[0]["columns"][0]["name"], myJson[0]["columns"][1]["name"]])

sql_query4 = ("""SELECT
              CASE
              WHEN shoredistance < 10 THEN 'Very Near Shore (<10km)'
              WHEN shoredistance BETWEEN 10 AND 50 THEN 'Near Shore (10–50km)'
              WHEN shoredistance BETWEEN 50 AND 200 THEN 'Offshore (50–200km)'
              ELSE 'Far Offshore (>200km)'
              END AS shore_band,
              brackish,
              COUNT(occurrenceID) AS num_occurrences
              FROM pelagis
              WHERE brackish != 'nan'
              GROUP BY 1, 2
              ORDER BY num_occurrences DESC""")

#Setup the URL to be used to invoke the GridDB WebAPI to retrieve data from the container
url = base_url + 'sql'

# Construct the request body
request_body = json.dumps([
    {
        "type": "sql-select",
        "stmt": sql_query4
    }
])

# Disable HTTP debugging
http.client.HTTPConnection.debuglevel = 0

#Invoke the GridDB WebAPI
data_req4 = requests.post(url, data=request_body, headers=header_obj)

#Process the response received and construct a Pandas dataframe with the data from the response
myJson = data_req4.json()
occurrences_by_shore_distance_df = pd.DataFrame(myJson[0]["results"], columns=[myJson[0]["columns"][0]["name"], myJson[0]["columns"][1]["name"], myJson[0]["columns"][2]["name"]])

sql_query5 = """SELECT
                IFNULL(p.vernacularName,'Not Specified') as VernacularName,
                AVG(p.sst) AS avg_sst,
                AVG(p.sss) AS avg_sss,
                COUNT(p.occurrenceID) AS num_occurrences
                FROM pelagis p
                WHERE p.genus IS NOT NULL
                GROUP BY p.vernacularName
                HAVING COUNT(p.occurrenceID) > 10
                ORDER BY num_occurrences DESC
                LIMIT 15
                """

#Setup the URL to be used to invoke the GridDB WebAPI to retrieve data from the container
url = base_url + 'sql'

# Construct the request body
request_body = json.dumps([
    {
        "type": "sql-select",
        "stmt": sql_query5
    }
])

# Disable HTTP debugging
http.client.HTTPConnection.debuglevel = 0

#Invoke the GridDB WebAPI
data_req5 = requests.post(url, data=request_body, headers=header_obj)

#Process the response received and construct a Pandas dataframe with the data from the response
myJson = data_req5.json()
avg_metrics_df = pd.DataFrame(myJson[0]["results"], columns=[myJson[0]["columns"][0]["name"], myJson[0]["columns"][1]["name"], myJson[0]["columns"][2]["name"], myJson[0]["columns"][3]["name"]])