from elasticsearch import Elasticsearch
from datetime import datetime
from dateutil.relativedelta import relativedelta

# Connect to your Elasticsearch instance
CLOUD_ID = "SCIC_ElasticSearch:dXMtY2VudHJhbDEuZ2NwLmNsb3VkLmVzLmlvOjQ0MyRmYWQyOTUwYTllNDE0MTE2YjIwN2QzYTE0MzQyMzcyMSRjZDI5OTY4MTQ1ODA0MDU4OWU4ZmVmMWQ2NDViYzNjMw=="

ELASTIC_PASSWORD = "8oKIqy312EBsAPzWT64NUzji"

es = Elasticsearch(cloud_id=CLOUD_ID, basic_auth=("elastic", ELASTIC_PASSWORD))
# Define the index name
index_name = "upload1"

# Function to get records for a specified month and year
def get_monthly_records(year, month, model_name):
    # Initialize a dictionary to hold results for the specified month
    monthly_results = {}

    # Calculate the start and end dates for the month
    start_date = datetime(year, month, 1)
    end_date = (start_date + relativedelta(months=1)).replace(day=1)

    # Convert dates to epoch milliseconds
    start_epoch = int(start_date.timestamp() * 1000)
    end_epoch = int(end_date.timestamp() * 1000)

    # Elasticsearch query for the specific month
    query = {
        "query": {
            "bool": {
                "must": [
                    {"term": {"modelname.keyword": model_name}},  # Exact match using 'term' query
                    {
                        "range": {
                            "created_at": {
                                "gte": start_epoch,
                                "lt": end_epoch  # use 'lt' for exclusive end date
                            }
                        }
                    }
                ]
            }
        }
    }

    # Execute the search
    response = es.search(index=index_name, body=query)

    # Store results in the dictionary
    monthly_results[start_date.strftime("%B %Y")] = response['hits']['hits']

    return monthly_results

# Example usage
# input_year = 2024  # Replace with the year you want to query
# input_month = 9    # Replace with the month you want to query (1 for January, 2 for February, etc.)
# results = get_monthly_records(input_year, input_month, "gpt-4-0125-preview")

# mon_gpt4prev = []
# # Print the results for the specified month
# for month, records in results.items():
#     print(f"Records for {month}:")
#     tom = 0
#     for record in records:
#         # print(record['_source'])
#         if 'human token' in record['_source']:
#             tom += record['_source']['human token']
#         elif 'assistant token' in record['_source']:
#             tom += record['_source']['assistant token']
#     mon_gpt4prev.append(tom)
#     print("\n")

    
input_year = 2024  # Replace with the year you want to query
input_month = 9    # Replace with the month you want to query (1 for January, 2 for February, etc.)
results = get_monthly_records(input_year, input_month, "gpt-4-0125-preview")

mon_gpt4prev = [0] * 31  # Initialize a list for 31 days
# Print the results for the specified month
for month, records in results.items():
    print(f"Records for {month}:",len(records))
    
    for record in records:
        # Check if the record contains a date field (assuming 'created_at' is a timestamp in milliseconds)
        timestamp_ms = record['_source'].get('created_at')  # Assuming each record has a 'created_at' field
        if timestamp_ms:  # Ensure the timestamp is present
            # Convert milliseconds to seconds for fromtimestamp
            timestamp_s = timestamp_ms / 1000.0
            date = datetime.fromtimestamp(timestamp_s)  # Convert to datetime
            day = date.day - 1  # Get the day (0-indexed for list)
            
            # Accumulate tokens for the specific day
            if 'human token' in record['_source']:
                mon_gpt4prev[day] += record['_source']['human token']
            elif 'assistant token' in record['_source']:
                mon_gpt4prev[day] += record['_source']['assistant token']
    
    # Print daily token counts for the month
    for day in range(1, 32):
        if day <= len(mon_gpt4prev):  # Ensure we only print valid days
            print(f"Day {day}: {mon_gpt4prev[day - 1]} tokens")
    print("\n")

# Optionally, trim the list to the number of days in the selected month
if input_month in [4, 6, 9, 11]:  # April, June, September, November have 30 days
    mon_gpt4prev = mon_gpt4prev[:30]
elif input_month == 2:  # February (consider leap year)
    if (input_year % 4 == 0 and input_year % 100 != 0) or (input_year % 400 == 0):
        mon_gpt4prev = mon_gpt4prev[:29]  # Leap year
    else:
        mon_gpt4prev = mon_gpt4prev[:28]  # Non-leap year

input_year = 2024  # Replace with the year you want to query
# results = get_monthly_records(input_year,"gpt-4o")
input_month = 9    # Replace with the month you want to query (1 for January, 2 for February, etc.)
results = get_monthly_records(input_year, input_month, "gpt-4o")
mon_gpt4=[]
# Print the results for each month
for month, records in results.items():
    print(f"Records for {month}:")
    tom=0
    for record in records:
        # print(record['_source'])
        if 'human token' in record['_source']:
            tom=tom+record['_source']['human token']
        elif 'assistant token' in record['_source']:
            tom=tom+record['_source']['assistant token']
    mon_gpt4.append(tom)
    print("\n")