
"""
-----------------------------------
# Instructions for Code Execution
-----------------------------------

## Overview
This script is designed to process and summarize large numbers of web-page files (in .txt format), focusing on extracting ESG content. 
It utilizes the OpenAI's GPT-3.5 Turbo model for generating ESG summaries and scores. 
To handle potential API rate limits or temporary issues, the script includes a backoff delay mechanism during requests, ensuring reliable data processing under varying network conditions and API response statuses.

## Folder Structure
- The script is stored in the 'Code' folder. The 'Temp Data', and  'Temp Output' are saved in the 'Data' folder. 'Data' and 'Results' folders are siblings to the 'Code' folder.
- **Data**: Contains the source files in the `Firmfiles_cleaned` folder.
- **Temp Data**: Used for intermediate processing (storing large files, split documents, and concatenated summaries).
- **Temp Output**: Stores intermediary CSV files generated during processing.
- **Results**: Contains the final output with summaries and ESG scores.

## Prerequisites
- An active OpenAI API key updated in the `.env` file.
- Required Python packages installed (run `pip install -r requirements.txt`).

## Processing Steps
1. **Token Counting**: The script counts tokens in each web-page file and saves the counts to a CSV file.
2. **File Segregation**: Files larger than 3500 tokens are copied to a designated folder for specialized processing.
3. **Summarization and Scoring (Small Files)**: Files with ≤ 3500 tokens are processed through the GPT-3.5 Turbo model to generate ESG summaries and scores.
4. **Splitting Large Files**: Larger files are split into smaller segments suitable for GPT-3.5 Turbo processing.
5. **Summarization and Scoring (Large Files)**: The split segments of larger files are individually summarized for their ESG content.
6. **Concatenation**: ESG Summaries of split segments are concatenated to form complete ESG summaries of the original large files.
7. **Final Scoring**: All concatenated ESG summaries undergo a final summarization and scoring phase using the GPT-3.5 Turbo model.

## Output
- The final csv with ESG summaries and scores is  generated and saved for all the files to 'summaries_with_scores.csv'in the "Results" folder.

## Note
- Ensure the `.env` file with the API key is correctly placed in the working directory.
- All necessary source data files are in the sub-folder "Firmfiles_cleaned" of 'Data' folder before execution.

"""

import os
import csv
import pandas as pd
import shutil
import nltk
nltk.download('punkt')
import openai
from dotenv import load_dotenv
import time
import tiktoken
#------------ Set up logging ------------ 
import log_config
log_config.setup_logging()
import logging
logging.info("Running script 'Main GPT Block.py'")


#------------ Load OpenAI API Key from the environment ------------
# Load environment variables from .env file
load_dotenv()
# Set API key in the environment variables (.env file)
api_key = os.getenv("OPENAI_API_KEY")
openai.api_key = api_key

#------------ Folder locations here ------------ 
# Define folder structure directly relative to the script's location
source_folder = '../Data/Firmfiles_cleaned'
temp_data_folder = '../Data/Temp Data'
temp_output_folder = '../Data/Temp Output'
results_folder = '../Results'

# Create necessary folders if they exist, delete them and recreate
def setup_folder(folder_path):
    if os.path.exists(folder_path):
        shutil.rmtree(folder_path)
    os.makedirs(folder_path)

# Set up each folder
setup_folder(temp_data_folder)
setup_folder(temp_output_folder)


# Define paths for generated files and folders
csv_file_path = os.path.join(temp_output_folder, 'token_count.csv')
destination_folder_large_files = os.path.join(temp_data_folder, 'greaterthan3500')
destination_folder_split = os.path.join(temp_data_folder, 'greaterthan3500_split')
Summaries_Scores = os.path.join(results_folder, 'summaries_with_scores.csv')
Summaries = os.path.join(temp_output_folder, 'summariesgreaterthan3500.csv')
Summaries_concatenated = os.path.join(temp_data_folder, 'greaterthan3500_concatenated')
csv_file_path2 = os.path.join(temp_output_folder, 'token_count_greaterthan3500_concatenated.csv')

# Create subdirectories within the main folders
setup_folder(destination_folder_large_files)
setup_folder(destination_folder_split)
setup_folder(Summaries_concatenated)
#------------ Below code first calculates the number of tokens present in each of the web-page files to a csv file ------------

# Function to count tokens
def count_tokens(string: str, encoding_name: str) -> int:
    encoding = tiktoken.get_encoding(encoding_name)
    num_tokens = len(encoding.encode(string))
    return num_tokens

# Function to count the tokens and save it in a csv file
def count_tokens_csv(folder_path, output_csv):
    token_counts = []
    for file_name in os.listdir(folder_path):
        file_path = os.path.join(folder_path, file_name)
        if os.path.isfile(file_path):
            with open(file_path, 'r', encoding='latin-1', errors='replace') as file:
                content = file.read()
                tokens = count_tokens(content, "cl100k_base")
                token_counts.append((file_name, tokens))
    with open(output_csv, 'w', newline='', encoding='utf-8') as file:
        writer = csv.writer(file)
        writer.writerow(['File Name', 'Token Count'])
        writer.writerows(token_counts)
    print("Token counts saved to the CSV file.")

count_tokens_csv(source_folder, csv_file_path)    

#------------ Below code copies web-page files greater than 3500 token size into another folder ------------

# Function to copy files with token size greater than 3500 to another folder
def copy_large_files(source_folder, destination_folder, csv_file_path):
    df = pd.read_csv(csv_file_path)
    for index, row in df.iterrows():
        file_name = row['File Name']
        token_count = row['Token Count']
        source_file = os.path.join(source_folder, file_name)
        if os.path.isfile(source_file) and token_count > 3500:
            destination_file = os.path.join(destination_folder, file_name)
            shutil.copy(source_file, destination_file)
            print(f'Copied file {file_name} to the greater than 3500 tokens folder.')
        else:
            print(f'File {file_name} will not be copied.')

copy_large_files(source_folder, destination_folder_large_files, csv_file_path)



#------------ Sending web-page files (<= 3500 tokens) to the GPT 3.5 Turbo model to generate ESG Summary and ESG scores  ------------

# Function for OpenAI requests with retry mechanism
def request_with_retry(prompt, max_retries=3):
    retries = 0
    while retries <= max_retries:
        try:
            response = openai.ChatCompletion.create(
                model='gpt-3.5-turbo',
                messages=[{"role": "user", "content": prompt}],
                temperature=0.0
            )
            if 'choices' in response:
                content = response['choices'][0]['message']['content']
                return content
            else:
                retries += 1
                print(f"Request failed. Retrying ({retries}/{max_retries})...")
                time.sleep(2 ** retries)  # Exponential backoff delay
        except openai.error.OpenAIError as e:
            print(f"Request failed with error: {e}. Retrying ({retries}/{max_retries})...")
            retries += 1
            time.sleep(2 ** retries)  # Exponential backoff delay
    print("Max retries exceeded.")
    return None

# Document summarization and scoring functions
def summarize_document(prompt):
    return request_with_retry(prompt)

def calculate_score(prompt):
    return request_with_retry(prompt)

# Function to save summaries with scores
def save_summaries_with_scores(summaries_with_scores):
    output_file = Summaries_Scores
    mode = 'a' if os.path.exists(output_file) else 'w'
    with open(output_file, mode, newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        if mode == 'w':
            writer.writerow(['File Name', 'Summary', 'Score'])
        writer.writerows(summaries_with_scores)

# Function to summarize and score documents in a folder
def summarize_and_score_documents_in_folder(folder_path, input_csv, batch_size=100, wait_time=30):
    summaries_with_scores = []
    with open(input_csv, newline='') as csvfile:
        reader = csv.reader(csvfile)
        next(reader)  # Skip header row
        count = 0
        for row in reader:
            file_name, token_count = row
            token_count = int(token_count)
            if token_count <= 3500:
                file_path = os.path.join(folder_path, file_name)
                with open(file_path, 'r') as file:
                    document = file.read()
                    prompt1 = f"""Your task is to generate a short summary related to ESG content present in the document. \
                      If you do not find any ESG Content then do not write anything. Summarize the document\
                      below, delimited by triple backticks. \n\n \
                    Document: ```{document}```"""
                    summary = summarize_document(prompt1)
                    prompt2 = f"""Your task is to critically analyze the following document and assign an ESG \
                    score based on the information contained within it. Each pillar \
                    - Environment, Social, and Governance - should be evaluated separately. Afterward, calculate the overall \
                    ESG score by averaging the three individual scores. Provide a detailed explanation to justify your scoring.\
                    Grade strictly: a score of 10 should represent perfection, and anything less should reflect\
                    deficiencies. The scoring format should be: \
                    Environment_{{Score}}_Social_{{Score}}_Governance_{{Score}}_TotalESGScore_\
                    {{Average Score}}_Explanation_{{Explanation}}. Make sure that you follow the format. Your answer should \
                    always separate the scores and explanation with an underscore.\
                    For example, if the document doesn't containany ESG-related information, assign a score of 0.0 to each\
                    pillar, and calculate the overall ESG score as 0.0 and give your response in the format:\
                    Environment_0.0_Social_0.0_Governance_0.0_TotalESGScore_0.0_Explanation_No ESG Content   \                 
                    Your responses should be scientific, rational, \
                    and consistent.  All scores should be in float format. \
                    Ensure the explanation is consistent for the same document across different iterations.\n\n\
                    This means if a certain document has been given a particular ESG score, say 7.0, \
                    the same document should always receive the same score in future assessments. Your calculation\
                    accuracy is of utmost importance - there is no room for errors in your assessment. Always keep the\
                    scores and explanation clearly separated using underscores and brackets respectively.\n\n\
                    Here is the document for your \
                    evaluation:\n\nDocument: ```{document}```"""
                    score = calculate_score(prompt2)
                    summaries_with_scores.append((file_name, summary, score))
                    
                    count += 1
                    if count % batch_size == 0:
                        save_summaries_with_scores(summaries_with_scores)
                        summaries_with_scores = []
                        print(f"Processed and saved {count} files. Waiting for {wait_time} seconds...")
                        time.sleep(wait_time)
    
    # Save any remaining summaries and scores
    if summaries_with_scores:
        save_summaries_with_scores(summaries_with_scores)
    
    

# Set the batch size and wait time
batch_size = 100
wait_time = 30

# Process and summarize the documents in the folder
summarize_and_score_documents_in_folder(source_folder, csv_file_path, batch_size, wait_time)
print(f"Summaries with scores have been saved to 'summaries_with_scores.csv' for files with token length <=3500.")

#------------ Sending web-page files (> 3500 tokens) to the GPT 3.5 Turbo model to generate ESG Summary and ESG scores  ------------
# Function to split documents (>=3500 tokens) using NLTK tokenizer
def split_document(document, max_tokens):
    tokenizer = nltk.data.load('tokenizers/punkt/english.pickle')
    sentences = tokenizer.tokenize(document)
    split_documents = []
    current_doc = ""
    current_tokens = 0
    for sentence in sentences:
        sentence_tokens = count_tokens(sentence, "cl100k_base")
        if current_tokens + sentence_tokens > max_tokens:
            split_documents.append(current_doc)
            current_doc = ""
            current_tokens = 0
        current_doc += sentence + " "
        current_tokens += sentence_tokens
    if current_doc:
        split_documents.append(current_doc)
    return split_documents

# Function to process the documents by splitting and saving them to a destination folder
def process_documents(source_folder, destination_folder, max_tokens=3500):
    if not os.path.exists(destination_folder):
        os.makedirs(destination_folder)
    
    files = os.listdir(source_folder)
    for file_name in files:
        if file_name.endswith(".txt"):
            file_path = os.path.join(source_folder, file_name)
            with open(file_path, 'r') as file_obj:
                document = file_obj.read()
                split_docs = split_document(document, max_tokens)
                file_name_without_ext = os.path.splitext(file_name)[0]
                for i, split_doc in enumerate(split_docs):
                    doc_name = f"{file_name_without_ext}_{i}.txt"
                    doc_path = os.path.join(destination_folder, doc_name)
                    with open(doc_path, 'w') as output_file:
                        output_file.write(split_doc)
    
    print("Documents have been split and saved in the destination folder for files > 3500 token size.")

process_documents(destination_folder_large_files, destination_folder_split, max_tokens=3500)

def save_summaries(summaries):
    output_file = Summaries
    mode = 'a' if os.path.exists(output_file) else 'w'
    with open(output_file, mode, newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        if mode == 'w':
            writer.writerow(['File Name', 'Summary'])
        writer.writerows(summaries)

def summarize_documents_in_folder(folder_path, batch_size=100, wait_time=30):
    summaries = []
    files = os.listdir(folder_path)
    count = 0
    for file_name in files:
        if file_name.endswith(".txt"):
            file_path = os.path.join(folder_path, file_name)
            with open(file_path, 'r') as file:
                document = file.read()
                prompt = f"""Your task is to generate a short summary related to ESG content present in the document. \
                If you do not find any ESG Content then do not write anything. Summarize the document\
                below, delimited by triple backticks. \n\n \
                Document: ```{document}```"""
                summary = summarize_document(prompt)
                summaries.append((file_name, summary))
                count += 1
                if count % batch_size == 0:
                    save_summaries(summaries)
                    summaries = []
                    print(f"Processed and saved {count} files. Waiting for {wait_time} seconds...")
                    time.sleep(wait_time)
    
    # Save any remaining summaries
    if summaries:
        save_summaries(summaries)
    
    print(f"Summaries have been saved to 'summariesgreaterthan3500.csv'.")

summarize_documents_in_folder(destination_folder_split, batch_size=100, wait_time=30)    

# Function to concatenate the ESG summaries from the split documents into web-page summaries
def concatenate_summaries(csv_file_path, destination_folder):
    df = pd.read_csv(csv_file_path, sep=',', encoding='utf-8')
    df['Base File Name'] = df['File Name'].apply(lambda x: '_'.join(x.split('_')[:-1]))
    df['Summary'] = df['Summary'].fillna('')
    grouped_df = df.groupby('Base File Name')['Summary'].apply(lambda x: ' '.join(x.astype(str))).reset_index()
    os.makedirs(destination_folder, exist_ok=True)
    for index, row in grouped_df.iterrows():
        new_file_name = row['Base File Name'] + '.txt'
        with open(os.path.join(destination_folder, new_file_name), 'w') as file:
            file.write(row['Summary'])
    print(f"Summaries have been concatenated to 'greaterthan3500_concatenated' folder.")

concatenate_summaries(Summaries, Summaries_concatenated)

# Count the tokens for the concatenated summaries of the web-pages
count_tokens_csv(Summaries_concatenated, csv_file_path2)

summarize_and_score_documents_in_folder(Summaries_concatenated, csv_file_path2, batch_size, wait_time)

print(f"Summaries with scores have been saved to 'summaries_with_scores.csv' for files with token length > 3500.")
print(f"Execution completed: Summaries with scores have been generated and saved for all the files to 'summaries_with_scores.csv' .")


#------------ Log completion of the script ------------
logging.info("Main GPT Block.py'")