from .student import Student
from .teacher import Teacher


MAX_LENGTH = 300


class StudentTeacher:
    """
    A class representing a system where a student and teacher interact, with the student generating 
    responses and the teacher evaluating those responses. The teacher uses feedback to train the student.

    Attributes
    ----------
    student : Student
        The student object, which uses a language model (LLM) to generate responses.
    teacher : Teacher
        The teacher object, which evaluates the student's responses.

    Methods
    -------
    generate_response(query)
        Generates a response from both the student and teacher LLMs for a given query.
    train_student(data)
        Fine-tunes the student model using provided training data.
    """

    def __init__(self, student_llm, teacher_llm):
        """
        Initializes the StudentTeacher system with a student and teacher.

        Parameters
        ----------
        student_llm : object
            The language model (LLM) used by the student for generating responses.
        teacher_llm : object
            The language model (LLM) used by the teacher to evaluate student responses.
        """
        self.student = Student(student_llm)
        self.teacher = Teacher(teacher_llm)

    def generate_response(self, query):
        """
        Generates a response from both the student and teacher LLMs based on a query. 
        The student generates a response which is then evaluated by the teacher.

        Parameters
        ----------
        query : str
            The query to be processed by the student and teacher.

        Returns
        -------
        tuple
            A tuple containing:
            - student_response : str
                The response generated by the student.
            - teacher_response : str
                The evaluation of the student response by the teacher.
        """
        student_response = self.student.generate_response(query)
        teacher_response = self.teacher.evaluate_response(
            student_response[:MAX_LENGTH])
        return student_response, teacher_response

    def train_student(self, data):
        """
        Fine-tunes the student model using the provided training data.

        Parameters
        ----------
        data : iterable
            The data used for fine-tuning the student's model.
        """
        self.student.fine_tune(data)
