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

# Data Gathering
import requests
import zipfile

#Encoding and Serialization
import base64
import json

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

#Visualization Libraries
import seaborn as sns
import matplotlib.pyplot as plt
import plotly.graph_objects as go
import plotly.express as px
import pycountry

# -*- coding: utf-8 -*-
"""
This script downloads the data sets for examples in data subfolder of examples folder.

Created on Fri Aug 12 10:19:50 2024

@author: U.S. Securities and Exchange Commission.
"""

def download_and_unzip(url, extract_to='.'):
    """
    Downloads a ZIP file from a URL and extracts its contents
    
    Parameters
    ----------
    url : str
        URL pointing to the ZIP file.
    extract_to : str, optional
        Directory path where the contents will be extracted. The default is 'data'.

    Returns
    -------
    None.

    """
    # Define headers with a User-Agent
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.36 (Contact: your-email@example.com)'
    }
    
    # Send a GET request to the URL
    print(f"Downloading ZIP file from {url} ...")
    response = requests.get(url, headers=headers)
    
    # Raise an exception if the download fails
    response.raise_for_status()
    
    # Create a ZipFile object from the bytes of the ZIP file
    zip_file = zipfile.ZipFile(BytesIO(response.content))
    
    # Extract the contents of the Zip file
    print(f"Extracting the contents to {extract_to}...")
    zip_file.extractall(path=extract_to)
    zip_file.close()
    print("Extraction complete.")

base_dir = os.getcwd()
dataset_path = os.path.join(base_dir, 'datasets', 'data', '2025q2')

def load_2025q2_data():
    sub = pd.read_csv(os.path.join(dataset_path, 'sub.txt'), sep='\t')
    num = pd.read_csv(os.path.join(dataset_path, 'num.txt'), sep='\t')
    tag = pd.read_csv(os.path.join(dataset_path, 'tag.txt'), sep='\t')
    return sub, num, tag.drop_duplicates(subset='tag')

# Load and join data for 2025q1
sub_df, num_df, tag_df = load_2025q2_data()

num_df['ddate'] = pd.to_datetime(num_df['ddate'].astype(str), format='%Y%m%d', errors='coerce')
key_tags = ['Revenues', 'NetIncomeLoss', 'Assets', 'Liabilities']
filtered_df = num_df[num_df['tag'].isin(key_tags)]
filtered_df = filtered_df[filtered_df['value'].notnull()]
filtered_df = filtered_df[filtered_df['value'].abs() < 1e12]  # sanity check
filtered_df = filtered_df[filtered_df['qtrs'] == 1]
filtered_df['value'] = filtered_df['value'].astype(float)
#filtered_df['ddate'] = pd.to_datetime(filtered_df['ddate'])

username = "XXXX" # Provide your GridDB credentials here
password =   "XXXX"
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 XXX","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
    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

data_obj = generate_griddb_data_obj(sub_df, container_name="sub_data")

#Set up the GridDB WebAPI URL
url = base_url + 'containers'

#Invoke the GridDB WebAPI with the headers and the request body
x = requests.post(url, json = data_obj, headers = header_obj)

data_obj = generate_griddb_data_obj(num_df, container_name="num_data")

#Set up the GridDB WebAPI URL
url = base_url + 'containers'

#Invoke the GridDB WebAPI with the headers and the request body
x = requests.post(url, json = data_obj, headers = header_obj)

data_obj = generate_griddb_data_obj(tag_df, container_name="tag_data")

#Set up the GridDB WebAPI URL
url = base_url + 'containers'

#Invoke the GridDB WebAPI with the headers and the request body
x = requests.post(url, json = data_obj, headers = header_obj)

# Parse the JSON string into a Python list of dictionaries
request_body_sub_data = sub_df.to_json(orient='values')

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

# Parse the JSON string into a Python list of dictionaries
request_body_num_data = filtered_df.to_json(orient='values')

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

# Parse the JSON string into a Python list of dictionaries
request_body_tag_data = tag_df.to_json(orient='values')

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

sql_query1 = ("""
            SELECT s.name, s.form, n.tag, n.value
            FROM num_data n
            JOIN sub_data s ON n.adsh = s.adsh
            WHERE n.tag = 'Revenues' 
            AND s.form IN ('10-K', '10-Q')
            ORDER BY 3 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()
revenue_valid_submissions = 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"]])

sql_query2 = ("""
            SELECT 
                s.adsh,
                s.name,
                s.form,
                r.ddate as rddate,
                SUM(case when r.tag = 'Revenues' then r.value else 0 end) as revenue,
                SUM(case when r.tag = 'NetIncomeLoss' then r.value else 0 end) as net_income
            FROM sub_data s
            JOIN num_data r ON s.adsh = r.adsh
            WHERE r.tag in ('Revenues', 'NetIncomeLoss')
              AND r.value IS NOT NULL
              AND s.form IN ('10-K', '10-Q')
              and r.ddate = (select max(n.ddate) from num_data n where n.adsh = s.adsh and n.tag in ('Revenues', 'NetIncomeLoss'))
            GROUP BY s.adsh, s.name, s.form, r.ddate
            ORDER BY 5 DESC
            LIMIT 5;
            """)

#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()
final_statement = 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"], myJson[0]["columns"][5]["name"]])

sql_query3 = ("""
            SELECT s.adsh,s.name, 
            MAX(n.ddate) AS latest_due_date, 
            s.filed, 
            s.accepted, n.segments,t.tlabel, s.form,
            s.fp,
            count(s.adsh) as num_submissions
            FROM sub_data s
            JOIN num_data n 
            ON s.adsh = n.adsh 
            AND n.tag = 'Revenues'
            JOIN tag_data t
            ON t.tag = n.tag
            GROUP BY s.adsh, s.name,
            s.filed, s.accepted,
            n.segments,
            t.tlabel,
            s.form,
            s.fp
            HAVING count(s.adsh) > 1
            ORDER BY 10 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()
submissions_count = 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"], myJson[0]["columns"][5]["name"],myJson[0]["columns"][6]["name"],myJson[0]["columns"][7]["name"],myJson[0]["columns"][8]["name"],myJson[0]["columns"][9]["name"]])

def extract_segment_keys(segment_str):
    if pd.isna(segment_str):
        return []
    return [pair.split('=')[0] for pair in segment_str.strip(';').split(';') if '=' in pair]

# Create a new column with list of keys
submissions_count['segment_keys'] = submissions_count['segments'].apply(extract_segment_keys)

# Explode the list into separate rows
exploded = submissions_count[['name', 'segment_keys']].explode('segment_keys')

# Drop NaNs and count unique keys per company
segment_summary = (
    exploded.dropna()
    .groupby('name')['segment_keys']
    .nunique()
    .reset_index(name='distinct_segments')
    .sort_values(by='distinct_segments', ascending=False)
)

sql_query4 = ("""
            SELECT a.adsh, a.name 
            FROM ( 
                SELECT s.adsh, s.name, SUM(n.value) revenue
                FROM sub_data s
                JOIN num_data n ON s.adsh = n.adsh
                WHERE n.tag IN ('Revenues')
                AND n.uom = 'USD'
                GROUP BY s.adsh, s.name
            ) a 
            WHERE a.revenue > 100000000
            INTERSECT
            SELECT b.adsh, b.name 
            FROM ( 
                SELECT s.adsh, s.name, sum(n.value) netincomeloss
                FROM sub_data s
                JOIN num_data n ON s.adsh = n.adsh
                WHERE n.tag = 'NetIncomeLoss'
                AND n.uom = 'USD'
                GROUP BY s.adsh, s.name
            ) b
            WHERE b.netincomeloss < -500000000
           
            """)

#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()
R_D_Companies_count = pd.DataFrame(myJson[0]["results"], columns=[myJson[0]["columns"][0]["name"], myJson[0]["columns"][1]["name"]])

sql_query5 = ("""SELECT 
                cityba, 
                stprba, 
                countryba, 
                status, 
                COUNT(DISTINCT name) AS company_count,
                AVG(netincomeloss) AS avg_net_income
                FROM (
                    SELECT s.adsh, s.name, s.cityba, 
                    s.stprba, s.countryba, 
                    n.value AS netincomeloss,
                    'Profit' AS status
                    FROM sub_data s
                    JOIN num_data n ON s.adsh = n.adsh
                    WHERE n.tag = 'NetIncomeLoss' AND n.value > 0
                    UNION
                    SELECT s.adsh, s.name, s.cityba, 
                    s.stprba, s.countryba, 
                    n.value AS netincomeloss,
                    'Loss' AS status
                    FROM sub_data s
                    JOIN num_data n ON s.adsh = n.adsh
                    WHERE n.tag = 'NetIncomeLoss' AND n.value < 0
                    ) AS profit_loss
                GROUP BY cityba, stprba, countryba, status;
""")

#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()
density_distribution = 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"],myJson[0]["columns"][5]["name"]])
density_distribution['log_income'] = density_distribution['avg_net_income'].apply(
    lambda x: np.sign(x) * np.log1p(abs(x))
)