# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_claude_corrector.ipynb.

# %% auto 0
__all__ = ['ClaudeDiarizationCorrector']

# %% ../nbs/00_claude_corrector.ipynb 6
from claudette import Chat, models
import os

class ClaudeDiarizationCorrector:
    def __init__(self, input_path: str, output_path: str, chunk_size: int = 20000):
        """
        Initialize the ClaudeDiarizationCorrector.
        
        Args:
            input_path (str): Path to the input diarization transcript file.
            output_path (str): Path where the corrected transcript will be saved.
            chunk_size (int, optional): Maximum size of each chunk for processing. Defaults to 20000.
        """
        self.input_path = input_path
        self.output_path = output_path
        self.chunk_size = chunk_size
        self.model = models[1]  # sonnet 3.5
        self.chat = Chat(self.model, sp="""You are a helpful and concise assistant that will correct a diarization script.""")

    def read_and_chunk(self):
        """
        Read the input file and split it into chunks.
        
        Returns:
            list: A list of text chunks.
        """
        with open(self.input_path, 'r', encoding='utf-8') as file:
            text = file.read()

        lines = text.split('\n')
        chunks = []
        current_chunk = ""

        for line in lines:
            if len(current_chunk + '\n' + line) > self.chunk_size:
                chunks.append(current_chunk)
                current_chunk = line
            else:
                if current_chunk:
                    current_chunk += '\n' + line
                else:
                    current_chunk = line

        if current_chunk:
            chunks.append(current_chunk)
        return chunks

    def process_chunk(self, chunk: str, index: int, total_chunks: int):
        """
        Process a single chunk of the diarization transcript.
        
        Args:
            chunk (str): The text chunk to process.
            index (int): The index of the current chunk.
            total_chunks (int): The total number of chunks.
        
        Returns:
            str: The corrected text chunk.
        """
        if index == 0:
            prompt = f'''In the speaker diarization transcript below, some words are potentially misplaced due to bleeding. As in one speaker is talking then immediately after a second speaker is talking, and the first words of the second speaker are attributed to the last words of the second speaker. I need you to read the meaning and context of the sentences and make your best guess about how the sentences should be re-distributed 
                      They may also be displaced in the following manner-- Two or more speakers speak, and the words are all attributed to one speaker, when actually 2 or three speakers spoke like this: speaker_1 "talking" speaker_2 "talking" speaker_1 "talking", but all the transcription captures is speaker_1 "talking talking talking" 
                      Please correct those words and move them to the right speaker. Directly show the corrected transcript without explaining what changes were made or why you made those changes. 
                      This is the first of {total_chunks} parts. Correct this chunk in isolation. Text to correct is indicated by this tag <text-to-correct></text-to-correct>
                      Feel free to attempt to correct the speaker identified, for example if it says it is speaker_01, but you think it is speaker_02, correct that, also if you know the actual name of a character, replace the "SPEAKER_01" with it's actual character name
                      <text-to-correct>{chunk}</text-to-correct>
                    '''
        else:
            if index % 2 == 0:
                self.chat = Chat(self.model, sp="""You are a helpful and concise assistant that will correct a diarization script.""")
            prompt = f'''In the speaker diarization transcript below, some words are potentially misplaced.
                      Please correct those words and move them to the right speaker. Directly show the corrected transcript without explaining what changes were made or why you made those changes.
                      Continue diarization from previous {index} parts, maintaining consistent speaker attribution and context. The diarization transcript to correct is indicated by this tag <text-to-correct></text-to-correct>
                      This is part {index + 1} of {total_chunks}. Correct the section indicated by <previous-corrected-text></previous-corrected-text> only, and directly show the corrected transcript without explaining what changes were made or why you made those changes.
                      Feel free to attempt to correct the speaker identified, for example if it says it is speaker_01, but you think it is speaker_02, correct that, also if you know the actual name of a character, replace the "SPEAKER_01" with it's actual character name
                      <text-to-correct>{chunk}</text-to-correct>
                    '''
        
        r = self.chat(prompt)
        print(f"Processed chunk {index + 1} of {total_chunks}")
        corrected_text = r.content
        return corrected_text

    def process_conversation(self):
        """
        Process the entire conversation, chunk by chunk.
        
        Returns:
            str: The fully corrected transcript.
        """
        chunks = self.read_and_chunk()
        results = []

        for index, chunk in enumerate(chunks):
            corrected_text = self.process_chunk(chunk, index, len(chunks))
            results.append(corrected_text)

        processed_text = '\n'.join(results)
        self.write_results_to_file(processed_text)
        return processed_text

    def write_results_to_file(self, processed_text: str):
        """
        Write the processed text to the output file.
        
        Args:
            processed_text (str): The fully corrected transcript.
        """
        with open(self.output_path, 'w', encoding='utf-8') as file:
            file.write(processed_text)
        print(f"Results written to {self.output_path}")
