from __future__ import print_function

try:
    input = raw_input
except NameError:
    pass

# imports compatible across Python versions
import os, sys, csv, datetime, tempfile

from schoolutils.config import user_config, user_calculators
from schoolutils.grading import db, validators
from schoolutils.reporting import reports

# TODO: abstract from specific institution
from schoolutils.institutions.ucberkeley import bspace

def require(attribute, callback, message):
    """Require that an object has an attribute before executing a method.
       attribute should be the attribute name.
       callback should be a method that will mutate the object to provide
         the attribute when it doesn't exist; it will be called with the object
         as an argument
       message should be a message to print to the user before callback is
         called"""
    def method_factory(f):
        def method(self, *args, **kwargs):
            attr = getattr(self, attribute, None)
            while not attr:
                print("")
                print(message)
                callback(self)
                attr = getattr(self, attribute, None)
            return f(self, *args, **kwargs)
        method.__doc__ = f.__doc__  # preserve docstring for actions_menu
        return method
    return method_factory

class BaseUI(object):
    def __init__(self, options=None):
        """Initialize grading program UI.
           options, if provided, should be an options structure produced by
             optparse
        """
        if options:
            self.cli_options = options
        else:
            self.cli_options = None

        self.semester = None
        self.year = None
        self.current_courses = []
        self.course_id = None
        self.assignment_id = None

        self.initial_database_setup()
        self.initial_course_setup()
        self.initial_assignment_setup()

    def get_config_option(self, option_name, validator, default=None):
        """Return the appropriate config value from CLI options or user config.
           option_name should be an attribute to look for on both the options
             object and the user_config module.  CLI options override user_config
             values.
           validator will be applied to the value.
           Returns the validated value, or default if option is not
           supplied by the user or the user-supplied value does not
           pass validation.
        """
        val = (getattr(self.cli_options, option_name, '') or
               getattr(user_config, option_name, ''))
        try:
            return validator(val)
        except ValueError:
            return default

    def initial_database_setup(self):
        "Set db_file and db_connection from user config and CLI options"
        self.db_file = self.get_config_option('gradedb_file', file_path)
        if self.db_file and os.path.exists(self.db_file):
            try:
                self.db_connection = db.connect(self.db_file, create=False)
            except db.ConnectionError:
                self.db_connection = None
        else:
            self.db_connection = None

    def initial_course_setup(self):
        """Set semester, year, current_courses, and course_id from user config
           and CLI options"""

        self.semester = self.get_config_option('current_semester',
                                               validators.semester)
        self.year = self.get_config_option('current_year', validators.year)
        self.current_courses = user_config.current_courses
        course_num = self.get_config_option('default_course',
                                            validators.course_number)

        if not (self.db_connection and self.semester and self.year):
            # don't bother looking for a course without a semester and year
            # (but try otherwise, because these might identify one uniquely)
            return

        try:
            self.course_id = db.ensure_unique(
                db.select_courses(self.db_connection,
                                  semester=self.semester,
                                  year=self.year,
                                  number=course_num))
        except (db.NoRecordsFound, db.MultipleRecordsFound):
            sys.stderr.write("Unable to locate a unique default course; "
                             "ignoring.\n")

    def initial_assignment_setup(self):
        "Set assignment_id using user config and CLI options"
        if not (self.db_connection and self.course_id):
            return

        if user_config.use_last_due_assignment:
            try:
                self.select_last_due_assignment()
            except AttributeError: # select_last_due_assignment is currently defined by SimpleUI
                sys.stderr.write("Ignoring use_last_due_assignment.\n")

class SimpleUI(BaseUI):
    """Manages a simple (command line) user interface.
    """
    # Helpers for printing data to stdout
    STUDENT_FORMAT = '{last_name}, {first_name} (SID: {sid})'
    COURSE_FORMAT = '{number} - {name} ({semester} {year})'
    ASSIGNMENT_FORMAT = '{name} (due {due_date})'
    GRADE_FORMAT = '' # TODO

    def course_formatter(self, course_row):
        "Format COURSE_FORMAT with course from db"
        return self.COURSE_FORMAT.format(year=course_row['year'],
                                         semester=course_row['semester'],
                                         number=course_row['number'],
                                         name=course_row['name'])

    def student_formatter(self, student_row):
        "Format STUDENT_FORMAT with student from db"
        return self.STUDENT_FORMAT.format(last_name=student_row['last_name'],
                                          first_name=student_row['first_name'],
                                          sid=student_row['sid'])

    def assignment_formatter(self, assignment_row):
        return self.ASSIGNMENT_FORMAT.format(name=assignment_row['name'],
                                             due_date=assignment_row['due_date'])

    def grade_formatter(self, grade_row):
        pass

    # BaseUI overrides
    def initial_database_setup(self):
        """Set db_file and db_connection from user config and CLI options.
           Query the user if automatic database connection fails."""
        super(SimpleUI, self).initial_database_setup()
        if self.db_connection:
            return
        if not self.db_file:
            return self.change_database()

        err_msg = ''
        if not os.path.exists(self.db_file):
            prompt = "No existing database at %s.\nCreate? (Y/N) " % self.db_file
            if typed_input(prompt, yn_bool):
                try:
                    self.db_connection = db.connect(self.db_file, create=True)
                except db.ConnectionError as e:
                    err_msg = ("FAILED to create database at {path}.\n"
                               "Error was: {err}".format(path=self.db_file, err=e))
        else:
            # retry automatic connection, mostly to get error message
            try:
                self.db_connection = db.connect(self.db_file, create=False)
            except db.ConnectionError as e:
                err_msg = ("FAILED to open file at {path} as a grade database.\n"
                           "Error was: {err}".format(path=self.db_file, err=e))

        if not self.db_connection:
            if err_msg:
                print(err_msg)
            print("Could not open a grade database based on your settings "
                  "in config.py.\n"
                  "Check the value of your gradedb_file setting and "
                  "permissions of your database \nfile and enclosing directory.")
            if typed_input("Enter database path manually? (Y/N) ", yn_bool):
                return self.change_database()
            else:
                # setup attempts failed, but we can still let the user
                # go to the main menu
                self.db_connection = None

    # Actions which can be @require-d:
    def close_database(self):
        """Close the current database connection."""
        if self.db_connection:
            print("Closing current database located at: %s" % self.db_file)
            self.db_connection.commit()
            self.db_connection.close()
            self.db_connection = None
            self.db_file = None
            # these fields will now be invalid, so erase them too:
            self.course_id = None
            self.assignment_id = None

    def change_database(self):
        """Open a new database.
           Closes the current database connection (if any) and opens another."""
        if self.db_connection:
            self.close_database()

        db_path = typed_input("Enter path to grade database: ", file_path)
        create = False
        if not os.path.exists(db_path):
            create = typed_input(
                "No existing database at %s.\nCreate (Y/N)? " % db_path,
                yn_bool)
        try:
            self.db_file = db_path
            self.db_connection = db.connect(db_path, create=create)
        except db.ConnectionError as e:
            print("Could not open {path} as a grade database.\n"
                  "Error was: {err}".format(path=db_path, err=e))

            if typed_input("Try again? (Y/N) ", yn_bool):
                return self.change_database()
            else:
                print("Database change aborted.")

    def get_student(self, create=False):
        """Lookup a student in the database, trying several methods.
           If create is True, allow (and offer) creating a new student using
             entered criteria if none exists.
           Returns student row.
        """
        student = None
        first_name = ''
        last_name = ''
        sid = ''
        email = ''

        # Helpers
        class UniqueFound(Exception):
            # dummy class for flow control
            pass

        def quit_if_unique(students):
            "Stop searching if a unique student has been located"
            if len(students) == 1:
                raise UniqueFound
            else:
                print("%d students found." % len(students))

        def students_menu(students):
            "Select a student (or None) from a menu"
            return self.options_menu(
                "Select a student:",
                students,
                self.student_formatter,
                allow_none=create)

        students = []
        try:
            print("Enter student data to lookup or create student. "
                  "Search uses fuzzy matching on name and email fields.\n"
                  "Use Ctrl-C to stop search and select from list.")
            sid = typed_input("Enter SID: ", validators.sid, default='')
            students = db.select_students(self.db_connection, sid=sid)
            quit_if_unique(students)

            last_name = typed_input("Enter last name: ", validators.name,
                                    default='')
            students = db.select_students(self.db_connection,
                                          sid=sid,
                                          last_name=last_name,
                                          fuzzy=True)
            quit_if_unique(students)

            first_name = typed_input("Enter first name: ", validators.name,
                                     default='')
            students = db.select_students(self.db_connection,
                                          sid=sid,
                                          last_name=last_name,
                                          first_name=first_name,
                                          fuzzy=True)
            quit_if_unique(students)

            email = typed_input("Enter email: ", validators.email, default='')
            students = db.select_students(self.db_connection,
                                          sid=sid,
                                          last_name=last_name,
                                          first_name=first_name,
                                          email=email,
                                          fuzzy=True)
            quit_if_unique(students)

        except UniqueFound:
            if create:
                # edge case: unique student found, but user may still
                # want to add a new one
                student = students_menu(students)
            else:
                student = students[0]

        except KeyboardInterrupt:
            # don't bother with a selection menu if no students met all
            # the search criteria at the time the user presses Ctrl-C
            if students:
                student = students_menu(students)
            elif not create:
                # if no students have been found and creation is not an option,
                # user wants to escape; so error should propagate
                raise KeyboardInterrupt

        if student:
            pass
        elif create and typed_input("No student found; create? (Y/N) ", yn_bool):
            # give user a chance to update values they provided and provide values
            # they didn't enter
            vals = {'last_name': last_name, 'first_name': first_name,
                    'sid': sid, 'email': email}
            print("Please provide data for the student to be created:")
            vals = self.edit_student_dict(from_dict=vals)
            vals.pop('student_id') # we're creating a new record
            student_id = db.create_student(self.db_connection, **vals)
            student = db.select_students(self.db_connection,
                                         student_id=student_id)[0]
        else:
            print("Could not locate student with these criteria; please try again.")
            return self.get_student(create=create)

        print("Selected: %s" % self.student_formatter(student))
        return student

    # Top-level actions:
    def main_loop(self):
        "Main menu"
        while True:
            self.print_db_info()
            self.print_course_info()
            self.print_assignment_info()
            self.actions_menu(
                "Main menu.",
                [self.change_database,
                 self.edit_courses,
                 self.edit_assignments,
                 self.import_students,
                 self.edit_student,
                 self.enter_grades,
                 self.edit_grades,
                 self.calculate_grades,
                 #self.import_grades,
                 self.export_grades,
                 self.grade_report,
                 self.exit])

            # commit after successful completion of any top-level action
            # to avoid data-loss
            if self.db_connection:
                self.db_connection.commit()

    @require('db_connection', change_database,
             "A database connection is required to edit courses.")
    def edit_courses(self):
        """Edit courses.
           Select, create, edit and delete courses."""
        def edit_course(c):
            make_default_clause = lambda v: (" (default: %s)" % str(v)
                                             if v else '')
            name_prompt = "Enter course name{default}: ".format(
                default=make_default_clause(c['name']))
            course_name = typed_input(name_prompt, validators.course_name,
                                      default=c['name'])
            year_prompt = "Enter year{default}: ".format(
                default=make_default_clause(c['year']))
            year = typed_input(year_prompt, validators.year, default=c['year'])
            sem_prompt = "Enter semester{default}: ".format(
                default=make_default_clause(c['semester']))
            semester = typed_input(sem_prompt, validators.semester,
                                   default=c['semester'])
            num_prompt = "Enter course number{default}: ".format(
                default=make_default_clause(c['number']))
            course_num = typed_input(num_prompt, validators.course_number,
                                     default=c['number'])

            course_id = db.create_or_update_course(
                self.db_connection,
                course_id=c['id'],
                year=year,
                semester=semester,
                name=course_name,
                number=course_num)

            return db.select_courses(self.db_connection, course_id=course_id)[0]

        create_course = lambda: edit_course({
                'id': None,
                'name': None,
                'year': None,
                'semester': None,
                'number': None,
                })

        def delete_course(c):
            existing_assignments = db.select_assignments(self.db_connection,
                                                         course_id=c['id'])
            existing_grades = [g for g in db.select_grades_for_course_members(
                                       self.db_connection,
                                       course_id=c['id'])
                               if g['value'] is not None]
            enrollees = db.select_students(self.db_connection, course_id=c['id'])
            if existing_assignments:
                print("WARNING: there are %d existing assignments for this course, "
                      "with %d associated grades.\n"
                      "Deleting this course will DELETE THESE ASSIGNMENTS AND GRADES, "
                      "and UNENROLL %d STUDENTS." %
                      (len(existing_assignments), len(existing_grades), len(enrollees)))
                if not typed_input("Delete anyway? (Y/N) ", yn_bool):
                    return False

            # deselect if this course was the currently selected course
            if self.course_id == c['id']:
                self.course_id = None
            if self.assignment_id in existing_assignments:
                self.assignment_id = None

            return db.delete_course_etc(self.db_connection,
                                        course_id=c['id'])

        def select_course(c):
            if self.course_id != c['id']:
                self.course_id = c['id']
                self.assignment_id = None
                print("Selected %s as current course."
                      % self.course_formatter(c))
            if user_config.use_last_due_assignment:
                self.select_last_due_assignment()
                if self.assignment_id:
                    assignment = db.select_assignments(
                        self.db_connection,
                        assignment_id=self.assignment_id)[0]
                    print("Selected %s as current assignment." %
                          self.assignment_formatter(assignment))

            print("")
            return True

        formatter = self.course_formatter
        header = "Courses"
        courses = db.select_courses(self.db_connection)

        self.edit_table(courses, header, formatter,
                        entity_type='course',
                        editor=edit_course,
                        creator=create_course,
                        deleter=delete_course,
                        selector=select_course)

    @require('db_connection', change_database,
             "A database connection is required to edit assignments.")
    @require('course_id', edit_courses,
             "A selected course is required to edit assignments.")
    def edit_assignments(self):
        """Edit assignments.
           Select, create, edit and delete assignments for the current course."""
        def edit_assignment(a):
            make_default_clause = lambda v: (" (default: %s)" % str(v)
                                             if v else '')
            name_prompt = "Enter assignment name{default}: ".format(
                default=make_default_clause(a['name']))
            name = typed_input(name_prompt, validators.assignment_name,
                               default=a['name'])
            desc_prompt = "Enter description{default}: ".format(
                default=make_default_clause(a['description']))
            description = typed_input(desc_prompt, str,
                                      default=a['description'])
            due_date_prompt = "Enter due date (YYYY-MM-DD){default}: ".format(
                default=make_default_clause(a['due_date']))
            due_date = typed_input(due_date_prompt, validators.date,
                                   default=a['due_date'])
            grade_type_prompt = "Enter grade type{default}: ".format(
                default=make_default_clause(a['grade_type']))
            grade_type = typed_input(grade_type_prompt, validators.grade_type,
                                     default=a['grade_type'])
            gt_unchanged = (grade_type == a['grade_type'])
            if grade_type == "points":
                wt_prompt = "Enter number of possible points{default}: ".format(
                    default=make_default_clause(
                        # don't use default if grade type changed
                        a['weight'] if gt_unchanged else ''))
            else:
                wt_prompt = "Enter grade weight (as decimal fraction of 1){default}: ".format(
                    default=make_default_clause(
                        # don't use default if grade type changed
                        a['weight'] if gt_unchanged else ''))
            weight = typed_input(wt_prompt, validators.grade_weight,
                                 default=a['weight'])

            a_id = db.create_or_update_assignment(
                self.db_connection,
                assignment_id=a['id'],
                course_id=self.course_id,
                name=name, description=description, grade_type=grade_type,
                due_date=due_date, weight=weight)

            return db.select_assignments(self.db_connection, assignment_id=a_id)[0]

        create_assignment = lambda: edit_assignment({
                'id': None,
                'name': None,
                'description': '',
                'due_date': None,
                'grade_type': None,
                'weight': None
                })

        def delete_assignment(a):
            existing_grades = db.select_grades(self.db_connection,
                                               assignment_id=a['id'])
            if existing_grades:
                print("WARNING: there are %d existing grades for this assignment.\n"
                      "Deleting this assignment WILL ALSO DELETE THEM." %
                      len(existing_grades))
                if not typed_input("Delete anyway? (Y/N) ", yn_bool):
                    return False

            # deselect if this assignment was the currently selected assignment
            if self.assignment_id == a['id']:
                self.assignment_id = None

            return db.delete_assignment_and_grades(self.db_connection,
                                                   assignment_id=a['id'])

        def select_assignment(a):
            self.assignment_id = a['id']
            print("Selected %s as current assignment.\n"
                  % self.assignment_formatter(a))
            return True

        format_str = ("{name: <20s} {due_date: <10s} {grade_type: <7s} {weight: <6} "
                      "{description: <32s}")
        formatter = lambda r: format_str.format(**r)
        header = format_str.format(name="Name", due_date="Due date",
                                   grade_type="Type", weight="Weight",
                                   description="Description")
        assignments = db.select_assignments(self.db_connection,
                                            course_id=self.course_id)

        self.edit_table(assignments, header, formatter,
                        entity_type='assignment',
                        editor=edit_assignment,
                        creator=create_assignment,
                        deleter=delete_assignment,
                        selector=select_assignment)

    @require('db_connection', change_database,
             "A database connection is required to select last due assignment.")
    @require('course_id', edit_courses,
             "A selected course is required to select last due assignment.")
    def select_last_due_assignment(self):
        """Select last due assignment.
           Selects the most recently due assignment in the current course, if any.
        """
        assignments = db.select_assignments(self.db_connection,
                                            course_id=self.course_id)
        most_recent = None
        current_date = datetime.date.today()
        past_assignments = [a for a in assignments
                            if a['due_date'] and
                               validators.date(a['due_date']) <= current_date]

        if not past_assignments:
            print("No assignments for this course have a due date in "
                  "the past; no current assignment is set.\n")
            self.assignment_id = None
            return

        most_recent = past_assignments[0]
        most_recent_date = validators.date(most_recent['due_date'])
        for a in past_assignments:
            due_date = validators.date(a['due_date'])
            if due_date > most_recent_date:
                most_recent = a
                most_recent_date = due_date

        self.assignment_id = most_recent['id']

    @require('db_connection', change_database,
             "A database connection is required to enter grades.")
    @require('course_id', edit_courses,
             "A selected course is required to enter grades.")
    @require('assignment_id', edit_assignments,
             "A selected assignment is required to enter grades.")
    def enter_grades(self):
        """Enter grades.
           Enter grades for the current assignment for individual students.
        """
        grade_type = db.select_assignments(self.db_connection,
                                           assignment_id=self.assignment_id)[0]['grade_type']
        grade_validator = validators.validator_for_grade_type(grade_type)

        print("")
        print("Use Control-C to finish entering grades.")
        while True:
            try:
                student = self.get_student()
                # avoid entering grades for non-member students:
                try:
                    membership = db.ensure_unique(db.select_course_memberships(
                            self.db_connection,
                            course_id=self.course_id,
                            student_id=student['id']))
                except db.NoRecordsFound:
                    course = db.select_courses(self.db_connection,
                                               course_id=self.course_id)[0]
                    print("{student} is not a member of {course}".format(
                            student=self.student_formatter(student),
                            course=self.course_formatter(course)))
                    # offer to add to course, but don't refuse to continue if not
                    if typed_input("Add this student to the course? (Y/N) ", yn_bool):
                        db.create_course_member(self.db_connection,
                                                student_id=student['id'],
                                                course_id=self.course_id)
                    else:
                        print("WARNING: grades for this student will not be "
                              "calculated or reported unless you add him or her "
                              "to the course later.")

                grade_id = None
                grade_val = typed_input("Enter grade value: ", grade_validator)
                existing_grades = db.select_grades(self.db_connection,
                                                   student_id=student['id'],
                                                   course_id=self.course_id,
                                                   assignment_id=self.assignment_id)
                if existing_grades:
                    print("Student has existing grades for this assignment.")
                    print("Existing grades are: %s" % ", ".join(
                            str(g['value']) for g in existing_grades))
                    update = typed_input("Update/overwrite? (Y/N) ", yn_bool)
                    if update:
                        if len(existing_grades) == 1:
                            print("Will update existing grade.")
                            grade = existing_grades[0]
                        else:
                            grade = self.options_menu(
                                "Select a grade to update.",
                                existing_grades,
                                # TODO: show timestamp?
                                lambda g: "{0}: {1}".format(
                                    g['assignment_name'],
                                    g['value']))
                        grade_id = grade['id']

                db.create_or_update_grade(self.db_connection,
                                          grade_id=grade_id,
                                          assignment_id=self.assignment_id,
                                          student_id=student['id'],
                                          value=grade_val)

            except KeyboardInterrupt:
                print("")
                break
            # TODO: shortcut here for changing to another assignment?
            # enter_grades_for_student method? (for a single student across all course assignments)

    @require('db_connection', change_database,
             "A database connection is required to edit grades.")
    @require('course_id', edit_courses,
             "A selected course is required to edit grades.")
    def edit_grades(self):
        """Edit grades.
           Edit a table of grades for the current course.
        """
        # TODO: handle case where student has multiple grades for an
        # assignment. Presently these are silently dropped in the
        # display, but when editing, user is asked to enter new grades
        # multiple times.
        def formatter(tbl_row):
            return row_fmt.format(name=self.student_formatter(tbl_row['student']),
                                  **dict((g['assignment_name'], str(g['value']))
                                         for g in tbl_row['grades']))

        def editor(tbl_row):
            print("Editing grades for %s." % self.student_formatter(tbl_row['student']))

            prompt = "New grade value for %s (default: %s): "
            any_updates = False
            for g in tbl_row['grades']:
                grade_validator = validators.validator_for_grade_type(g['grade_type'])
                old_val = g['value']
                # TODO: user can get stuck in a loop here if there is no existing
                # grade value but user doesn't want to enter one, because typed_input
                # will requery when default is None.  WONTFIX for now, because this
                # is an edge case, because it is easily (and correctly) escaped via
                # Control-C, and because generally, user should just enter the grade.
                new_val = typed_input(prompt % (g['assignment_name'], old_val),
                                      grade_validator, default=old_val)
                if new_val == old_val:
                    continue
                elif new_val:
                    any_updates = True
                    new_grade_id = db.create_or_update_grade(
                        self.db_connection,
                        student_id=g['student_id'],
                        assignment_id=g['assignment_id'],
                        grade_id=g['grade_id'], # may be None if grade didn't exist
                        value=new_val)

            if any_updates:
                return {
                    'student': tbl_row['student'],
                    'grades': db.select_grades_for_course_members(
                        self.db_connection,
                        course_id=self.course_id,
                        student_id=tbl_row['student']['id'])
                    }
            else:
                return tbl_row

        students = db.select_students(self.db_connection,
                                      course_id=self.course_id)
        all_grades = db.select_grades_for_course_members(
            self.db_connection,
            course_id=self.course_id)
        rows = [
            {'student': s,
             'grades': [r for r in all_grades if r['student_id'] == s['id']]}
            for s in students]

        # we use select_assignments here because it orders the
        # assignments by due date
        assignments = db.select_assignments(self.db_connection,
                                            course_id=self.course_id)
        assignment_names = [a['name'] for a in assignments]
        row_fmt = "{name: <40s}  "
        for a in assignment_names:
            row_fmt += "{" + a + ": <10s}"

        header = row_fmt.format(name="Student",
                                **dict((a,a) for a in assignment_names))

        self.edit_table(rows, header, formatter, editor=editor,
                        entity_type='grades for student')
        print("Grades updated successfully.")

    @require('db_connection', change_database,
             "A database connection is required to import students.")
    @require('course_id', edit_courses,
             "A selected course is required to import students.")
    def import_students(self):
        """Import students.
           Import students from a CSV file as members of the current course.
        """
        csv_path = typed_input("Enter path to CSV file with student data to import: ", file_path)
        if not os.path.exists(csv_path):
            retry = typed_input("File %s does not exist; try again? (Y/N) " %
                                csv_path, yn_bool)
            if retry:
                return self.import_students()
            else:
                return None

        # assume, for now, that CSV files have been generated by bSpace:
        students = bspace.roster_csv_to_students(csv_path)
        formatter = lambda s: "{last_name: <15s} {first_name: <20s} {email: <30s} {sid: <8}".format(**s)
        editor = lambda s: self.edit_student_dict(from_dict=s)
        creator = self.edit_student_dict
        header = formatter({'last_name': "Last name", 'first_name': "First name",
                            'email': "Email", 'sid': "SID"})
        students = self.edit_table(students, header, formatter,
                                   editor=editor, creator=creator,
                                   deleter=lambda s: True)

        for s in students:
            try:
                s['student_id'] = db.get_student_id(self.db_connection,
                                                    sid=s['sid'],
                                                    first_name=s['first_name'],
                                                    last_name=s['last_name'])
            except db.NoRecordsFound:
                s['student_id'] = None

            student_id = db.create_or_update_student(
                self.db_connection,
                sid=s['sid'],
                first_name=s['first_name'],
                last_name=s['last_name'],
                email=s['email'])
            course_member_id = db.create_course_member(
                self.db_connection,
                student_id=student_id,
                course_id=self.course_id)

        print("%d students imported successfully." % len(students))

    @require('db_connection', change_database,
             "A database connection is required to edit students.")
    def edit_student(self):
        """Add or edit students.
           Lookup students and modify their contact data and course memberships.
        """
        self.actions_menu(
            "What do you want to do?",
            [self.edit_student_info,
             self.edit_student_courses,
             self.edit_course_members])

    @require('db_connection', change_database,
             "A database connection is required to edit student information.")
    def edit_student_info(self):
        """Edit student contact data.
           Change name, SID, email, etc. for a single student.
        """
        # ugly hack: check result of get_student against list of existing students
        # to avoid asking user to enter student data 3(!) times (to search, create, edit)
        all_student_ids = [s['id'] for s in db.select_students(self.db_connection)]
        student = self.get_student(create=True)
        if student['id'] in all_student_ids:
            # this was an existing student, so run the editor:
            d = self.edit_student_dict(from_row=student)
            student_id = db.create_or_update_student(self.db_connection, **d)
        else:
            # this student was just created, so skip the editor;
            # get_student already gave user a chance to confirm data
            pass

        print("Student information updated.")

    @require('db_connection', change_database,
             "A database connection is required to edit course memberships.")
    def edit_student_courses(self):
        """Edit student courses.
           Add or remove a single student from one or more courses.
        """
        student = self.get_student()
        def add_to_course():
            all_courses = db.select_courses(self.db_connection)
            current_courses = db.select_courses(self.db_connection,
                                                student_id=student['id'])
            options = [c for c in all_courses if c not in current_courses]
            if not options:
                print("\nThis student is already enrolled in every course.")
                return None
            course = self.options_menu(
                "Which course should the student be added to?",
                options,
                self.course_formatter,
                allow_none=True)
            if course:
                course_id = course['id']
                db.create_course_member(self.db_connection,
                                        student_id=student['id'],
                                        course_id=course_id)
                print("Student added to %s" % self.course_formatter(course))
            else:
                print("Student not added to course.")

            return course

        def remove_from_course(course):
            course_id = course['id']
            db.delete_course_member(self.db_connection,
                                    student_id=student['id'],
                                    course_id=course_id)
            print("Student deleted from %s" % self.course_formatter(course))
            return True

        current_courses = db.select_courses(self.db_connection,
                                            student_id=student['id'])
        self.edit_table(
            current_courses,
            "Current courses for %s" % self.student_formatter(student),
            self.course_formatter,
            creator=add_to_course,
            deleter=remove_from_course)

        print("Student course memberships updated.")

    @require('db_connection', change_database,
             "A database connection is required to edit students in course.")
    @require('course_id', edit_courses,
             "A selected course is required to edit students in course.")
    def edit_course_members(self):
        """Edit course enrollments.
           Add or remove one or more students in the current course.
        """
        def add_to_course():
            student = self.get_student(create=True)
            db.create_course_member(self.db_connection,
                                    student_id=student['id'],
                                    course_id=self.course_id)
            print("Added %s to course." %
                  self.student_formatter(student))
            return student

        def remove_from_course(student):
            student_id = student['id']
            db.delete_course_member(self.db_connection,
                                    student_id=student_id,
                                    course_id=self.course_id)
            print("Deleted student %s from course." %
                  self.student_formatter(student))
            return True

        current_students = db.select_students(self.db_connection,
                                              course_id=self.course_id)
        current_course = db.select_courses(self.db_connection,
                                           course_id=self.course_id)[0]
        self.edit_table(
            current_students,
            "Current students in %s" % self.course_formatter(current_course),
            self.student_formatter,
            creator=add_to_course,
            deleter=remove_from_course,
            entity_type="student's membership")
        print("Course enrollments updated.")

    def import_grades(self):
        pass

    @require('db_connection', change_database,
             "A database connection is required to export grades.")
    @require('course_id', edit_courses,
             "A selected course is required to export grades.")
    def export_grades(self):
        """Export grades.
           Export grades for the current course to a CSV file.
        """
        out_file_name = typed_input("Enter a path to CSV file to export grades: ", file_path)
        if os.path.exists(out_file_name):
            print("Warning: file %s exists." % out_file_name)
            overwrite = typed_input("Overwrite? (Y/N) ", yn_bool)
            if not overwrite:
                print("Abort.")
                return

        out_file = open(out_file_name, 'w')

        # format, for now:
        # last_name + first_name, sid, grade1, grade2, grade3...
        # we use select_assignments here because it orders the
        # assignments by due date
        assignments = db.select_assignments(self.db_connection,
                                            course_id=self.course_id)
        assignment_names = [a['name'] for a in assignments]
        header = ["Name", "SID"] + assignment_names
        writer = csv.DictWriter(out_file, header)

        students = db.select_students(self.db_connection,
                                      course_id=self.course_id)

        # writeheader() became available in Python 2.7:
        try:
            writer.writeheader()
        except AttributeError:
            writer.writerow(dict(zip(header,header)))

        all_grades = db.select_grades_for_course_members(
            self.db_connection,
            course_id=self.course_id)

        for s in students:
            row = {}
            row["Name"] = "%s, %s" % (s['last_name'], s['first_name'])
            row["SID"] = s['sid']
            grades = [r for r in all_grades if r['student_id'] == s['id']]
            for g in grades:
                assignment_name = g['assignment_name']
                if assignment_name not in row:
                    row[assignment_name] = g['value']
                else:
                    print("Warning: multiple grades found for student %s "
                          "for assignment %s; only exporting first result."
                          % (self.student_formatter(s), assignment_name))
                    continue

            try:
                writer.writerow(row)
            except IOError:
                print("Warning: could not write row to CSV: %r." % row)
                continue

        out_file.close()
        print("Grades exported successfully to: %s.\n" % out_file_name)

    @require('db_connection', change_database,
             "A database connection is required to calculate grades.")
    @require('course_id', edit_courses,
             "A selected course is required to calculate grades.")
    def calculate_grades(self):
        """Calculate grades.
           Run the (user-defined) grade calculation function for students in the
           current course.
        """
        def save_calculated_grade(student_id,
                                  name='', # required, unless grade_id or assignment_id given
                                  value='', # required
                                  # to update an existing grade:
                                  grade_id=None,
                                  # to add a grade for an existing assignment:
                                  assignment_id=None,
                                  # to create a new assignment:
                                  description='(Assignment for calculated grade)',
                                  due_date=datetime.date.today(),
                                  grade_type=None,
                                  weight='CALC'):
            if not name and not (assignment_id or grade_id):
                raise ValueError("No assignment name given for calculated grade.")
            if value is None: # missing values not allowed, but 0/False/etc. OK
                raise ValueError("No value given for calculated grade %s." % name)

            if grade_id:
                db.update_grade(self.db_connection,
                                student_id=student_id,
                                grade_id=grade_id,
                                value=value)
                return grade_id

            if not assignment_id:
                try:
                    assignment_id = db.ensure_unique(db.select_assignments(
                        self.db_connection,
                        course_id=self.course_id,
                        name=name))
                except db.NoRecordsFound:
                    assignment_id = db.create_assignment(
                        self.db_connection,
                        course_id=self.course_id,
                        name=name,
                        description=description,
                        grade_type=grade_type,
                        due_date=due_date,
                        weight=weight)
                # MultipleRecordsFound should propagate

            # avoid storing calculated grades multiple times
            try:
                grade_id = db.ensure_unique(
                    db.select_grades(self.db_connection,
                                     student_id=student_id,
                                     assignment_id=assignment_id))
            except db.NoRecordsFound:
                grade_id = None

            row_id = db.create_or_update_grade(
                self.db_connection,
                student_id=student_id,
                assignment_id=assignment_id,
                grade_id=grade_id,
                value=value)

            return row_id

        course = db.select_courses(self.db_connection,
                                   course_id=self.course_id)[0]
        safe_num = course['number'].replace('-', '_').replace('.', '_')
        calc_name = ('calculate_grade_' + safe_num + '_' +
                     course['semester'].lower() + str(course['year']))
        calc_func = getattr(user_calculators, calc_name, None)

        if not calc_func:
            print("Could not locate grade calculation function %s. "
                  "Have you written it?" % calc_name)
            print("")
            return

        students = db.select_students(self.db_connection,
                                      course_id=self.course_id)
        all_grades = db.select_grades_for_course_members(
                self.db_connection,
                course_id=self.course_id)

        for s in students:
            grades = [r for r in all_grades
                      if (r['student_id'] == s['id'] and r['weight'] != 'CALC')]

            try:
                calculated_grades = calc_func(grades)
            except Exception as e:
                print("Failed to calculate grades for %s. "
                      "Error was: %s.  Skipping..." %
                      (self.student_formatter(s), e))
                continue
            if type(calculated_grades) is dict:
                # transpose to the list format to use save_calculated_grades
                calculated_grades = [
                    {'name': k, 'value': v}
                    for k, v in calculated_grades.items()]

            for cg in calculated_grades:
                save_calculated_grade(s['id'], **cg)

        print("Grade calculations ran successfully.\n")

    @require('db_connection', change_database,
             "A database connection is required to view a grade report.")
    @require('course_id', edit_courses,
             "A selected course is required to view a grade report.")
    def grade_report(self):
        """View grade report.
           See a report on grades in the current course."""
        r = reports.GradeReport(self.db_connection, course_id=self.course_id)
        r.run()
        print(r.as_text(compact=True))
        if typed_input("See and save the full report? (Y/N): ", yn_bool):
            # TODO: support for pager program?
            full_report = r.as_text(compact=False)
            print(full_report)

            course = db.select_courses(self.db_connection,
                                       course_id=self.course_id)[0]
            name = "grade_report_{number}_{semester}_{year}-".format(**course)
            with tempfile.NamedTemporaryFile(prefix=name, delete=False) as t:
                try:
                    # 3.x
                    report_bytes = full_report.encode("utf-8")
                except AttributeError:
                    # 2.x
                    report_bytes = full_report
                t.write(report_bytes)
                print("Full report saved at: %s\n" % t.name)

    def exit(self):
        """Quit grader.
           Closes database connection and exits."""
        self.close_database()
        exit(0)

    # Helper methods:
    def actions_menu(self, query, actions, default=None, escape=None):
        """Present a menu of actions to the user.
           query should be a string to print before the list of actions.
           actions should be a sequence callables.
             The docstring of each action is used to provide its menu entry.
           default, if provided, should be a callable to call if the user
             makes no selection
           escape, if provided, should be a callable that returns the user
             to a previous menu.
           Returns the return value of the selected action.
        """
        # TODO: add support for arguments to actions?

        if escape:
            actions = list(actions) + [escape]

        menu_format = "{0:>3d}: {1: <30} {2: <40}"
        short_descs = []
        long_descs = []
        for a in actions:
            lines = a.__doc__.splitlines()
            short_descs.append(lines[0])
            long_descs.append(' '.join([l.strip() for l in lines[1:]]))

        print("")
        print(query)
        print("Please select one of the following actions:")
        for i, a in enumerate(actions):
            print(menu_format.format(i, short_descs[i], long_descs[i]))

        if default is not None:
            prompt = "Which action? (default %d): " % actions.index(default)
        else:
            prompt = "Which action? (enter a number): "

        validator = lambda s: validators.int_in_range(s, 0, len(actions))
        i = typed_input(prompt, validator,
                        default=actions.index(default) if default in actions else None)
        print("") # visually separate menu and selection

        return actions[i]()

    def options_menu(self, query, options, formatter,
                     escape=None, allow_none=False):
        """Present a menu of options to the user.
           query should be a string to print before the list of options
           options should be a list of values
           formatter should be a callable that returns a string
             representation of a given option.
           escape, if provided, should be a callable to run in lieu of returning
             a selected value
           allow_none, if true, will add an option for the user to make no selection.
           Returns the selected option or None.
        """
        print("")
        print(query)
        menu_format = "{0:>3d}: {1}"
        max_index = len(options) - 1
        escape_option = None
        none_option = None
        for i, o in enumerate(options):
            print(menu_format.format(i, formatter(o)))
        if escape:
            escape_option = max_index
            max_index += 1
            print(menu_format.format(max_index, escape.__doc__.splitlines()[0]))
        if allow_none:
            # TODO: something wrong with logic here; escape doesn't work
            none_option = max_index
            max_index += 1
            print(menu_format.format(max_index, "None of the above"))

        validator = lambda s: validators.int_in_range(s, 0, max_index+1)

        idx = typed_input("Which option? (enter a number): ", validator)
        print("") # visually separate menu and selection input

        if idx < len(options):
            return options[idx]
        elif idx == escape_option:
            return escape()
        elif idx == none_option:
            return None

    def edit_table(self, rows, header, formatter, editor=None,
                   creator=None, deleter=None, selector=None, entity_type='row'):
        """Present a simple interface for reviewing and editing tabular data.
           rows should be a sequence of values for the user to review and edit
           header should be a string to print above the table
           formatter should be a function which, given a row value, returns a
             string representing the row for display to the screen
           editor, if provided, should be a function which, given a row value,
             returns a new row value based on user input
           creator, if provided, should be a function with no arguments which
             returns a new row value based on user input
           deleter, if provided, should be a function to call for its side
             effects with the row value when the user deletes a row.
             It should return False if deletion was aborted or unsuccessful,
             and a true value otherwise.
           selector, if provided, should be a function to call for its side
             effects with the row value when the user selects a row to be the
             currently-selected value (e.g., current assignment or course).
             It should return True if this function should exit after a selection
             is made, False otherwise.
           entity_type, if provided, will be used instead of the generic name 'row'
             when displaying the prompts for selecting, editing, creating, etc.
           Returns the edited rows.
        """
        editable_rows = [r for r in rows]
        header_underline = "-".ljust(80, "-")
        row_format = "{index: >5}: {frow: <73s}"
        def validator(s):
            s = s.strip().lower()
            if s.startswith('d'):
                action = 'd'
                idx = validators.int_in_range(s[1:], 0, len(editable_rows)+1)
            elif s.startswith('e'):
                action = 'e'
                idx = validators.int_in_range(s[1:], 0, len(editable_rows)+1)
            elif s.startswith('c'):
                action = 'c'
                idx = None
            else:
                action = 's'
                idx = validators.int_in_range(s, 0, len(editable_rows)+1)
            return action, idx

        while True:
            try:
                print("")
                print(row_format.format(index='Row #', frow=header))
                print(header_underline)
                if editable_rows:
                    for i, r in enumerate(editable_rows):
                        print(row_format.format(index=i, frow=formatter(r)))
                else:
                    print("(No %s data yet.)" % entity_type)

                print("")
                prompt = "Press Ctrl-C to end.\n"
                if creator:
                    prompt += "Enter 'c' to create a new %s. " % entity_type
                if editor:
                    prompt += "Prefix row # with 'e' to edit %s. " % entity_type
                if deleter:
                    prompt += "Prefix row # with 'd' to delete %s. " % entity_type
                if selector:
                    prompt += ("\nEnter row # to select as current %s. " %
                               entity_type)
                prompt += "\nWhat do you want to do? "

                action, idx = typed_input(prompt, validator)
                if action == 'e' and editor:
                    new_row = editor(editable_rows[idx])
                    if new_row: # editor might return None
                        editable_rows[idx] = new_row
                elif action == 'd' and deleter:
                    success = deleter(editable_rows[idx])
                    if success:
                        editable_rows.pop(idx)
                    else:
                        print("Deletion unsuccessful.")
                elif action == 'c' and creator:
                    new_row = creator()
                    if new_row: # creator might return None
                        editable_rows.append(new_row)
                elif action == 's' and selector:
                    should_exit = selector(editable_rows[idx])
                    if should_exit:
                        return editable_rows

            except KeyboardInterrupt:
                print("")
                break

        return editable_rows

    def edit_dict(self, d, validators=None, skip=None):
        """Simple interface for editing the values in a dictionary.
           Queries the user for a new value for each key in d, then
             updates the dictionary with a new validated value.
           d should be a dictionary of values with strings as keys
             The keys will be pretty printed, split along '_', when
             querying the user for new values.
           validators, if given, should be a dictionary mapping keys
             of d to functions validating input; str() will be used
             as a default validator if none is provided for a given key
           skip, if given, should be a list of keys in d to skip
             querying the user for
           Returns the updated dictionary.
        """
        if not validators:
            validators = {}
        if not skip:
            skip = []

        for k, v in d.items():
            if k in skip:
                continue

            validator = validators.get(k, str)
            pretty_field = ' '.join([word.strip().lower()
                                     for word in k.split('_')])
            pretty_current = (" (default: %s): " % v) if v else ": "
            new_v = typed_input("Enter %s%s" % (pretty_field, pretty_current),
                                validator, default=v)
            if new_v:
                d[k] = new_v

        return d

    def edit_student_dict(self, from_dict=None, from_row=None):
        """Convenience wrapper for edit_dict for student data.
           from_row, if provided, should be a row from the students table,
             with at least the following columns:
             id, last_name, first_name, sid, email
           from_dict, if provided, should have keys:
             last_name, first_name, sid, email
           You may not pass both from_row and from_dict; but you may pass
             neither.  In that case, the user will have to enter all fields.
           Returns a dictionary suitable to pass to *_student database
             functions as kwargs, with keys:
             student_id, last_name, first_name, sid, email
        """
        if from_dict and from_row:
            raise ValueError("You may not pass both from_dict and from_row")

        db_fields = ['student_id', 'last_name', 'first_name', 'sid', 'email']
        vlds = {'last_name': validators.name, 'first_name': validators.name,
                'sid': validators.sid, 'email': validators.email}
        d = {}
        for f in db_fields:
            if from_row:
                # student_id field comes from db as 'id' column
                d[f] = from_row[f if f != 'student_id' else 'id']
            elif from_dict:
                d[f] = from_dict.get(f, None)
            else:
                d[f] = None

        return self.edit_dict(d, validators=vlds, skip=['student_id'])

    def print_course_info(self):
        "Prints information about the currently selected course"
        if self.course_id:
            course = db.select_courses(self.db_connection, course_id=self.course_id)[0]
            print("Current course is: %s" % self.course_formatter(course))
        else:
            print("No course selected")

    def print_assignment_info(self):
        "Prints information about the currently selected assignment"
        if self.assignment_id:
            assignment = db.select_assignments(self.db_connection,
                                               assignment_id=self.assignment_id)[0]
            print("Current assignment is: %s" % self.assignment_formatter(assignment))
        else:
            print("No assignment selected")

    def print_db_info(self):
         "Prints information about the database connection"
         if self.db_connection:
             print("Database open at %s" % self.db_file)
         else:
             print("No current database connection")

#
# Utilities
#
def typed_input(prompt1, constructor, prompt2=None, default=None):
    """Get input and convert it to a given type.
       prompt1 should be an initial prompt for the user
       constructor should be a callable that converts a string
         to a value of the desired type or raises ValueError
       prompt2, if provided, should be a second prompt to use when
         the user provides an unacceptable value
       default, if provided, should be a default value to use when the
         user provides no input
    """
    if not prompt2:
        prompt2 = prompt1

    prompt = prompt1
    val = None
    while val is None:
        try:
            s = input(prompt)
            if not s and default is not None:
                val = default
            else:
                val = constructor(s)
        except ValueError:
            prompt = prompt2
            continue

    return val

#
# Constructors/validators
#
def file_path(s):
    """Convert a string to a file path.
       The passed string may contain '~' and will be expanded to an absolute
       path."""
    fp = s.strip()
    if not fp:
        raise ValueError("File path may not be empty")
    return os.path.abspath(os.path.expanduser(fp))

def yn_bool(s):
    """Convert a yes/no string to a boolean.
       Strings beginning with 'Y' and 'y' return True, with 'N' and 'n' return False."""
    b = s and s.strip().upper()[0]
    if b == 'Y':
        return True
    elif b == 'N':
        return False
    else:
        raise ValueError("Expected 'Y' or 'N': %s" % s)


from brownie import *
from dotenv import load_dotenv, find_dotenv
from brownie.convert.normalize import format_input
from brownie.convert import to_bytes
from utils.deploy_helpers import deploy_proxy, deploy_admin
import os
import json

def main():
    load_dotenv(dotenv_path="./.env", override=True)
    load_dotenv(find_dotenv())
    user = accounts.add(os.getenv("USER_PRIVATE_KEY"))

    f = open('./scripts/deploy_script/deploy_info.json')
    data = json.load(f)
    tokens = data["tokens"]
    deployed_contracts_addresses = data["deployedContracts"]

    # Getting instances of contracts
    usdc = ERC20PresetMinterPauserMock.at(tokens["USDC"])
    usdc_market = MErc20.at(deployed_contracts_addresses["UsdcMarketToken"])
    controller = Controller.at(deployed_contracts_addresses["Controller"])

    # Set amount of USDC to deposit(currently 5 USDC)
    amount_to_deposit = 5e6

    # Approving USDC
    usdc.approve(usdc_market, amount_to_deposit, {"from": user})

    # Depositing USDC tokens to USDC Market Token
    print("Your balance of USDC before depositing: ", usdc.balanceOf(user) / 1e6)
    print("Your balance of USDC Market Token before depositing:", usdc_market.balanceOf(user) / 1e8)

    usdc_market.mint(amount_to_deposit, {"from": user})

    print("Your balance of USDC after depositing: ", usdc.balanceOf(user) / 1e6)
    print("Your balance of USDC Market Token after depositing:", usdc_market.balanceOf(user) / 1e8)

    # Entering the market, marking that asset can be used as collateral
    controller.enterMarkets([usdc_market], {"from": user})


# !/usr/bin/env python3
"""
==== No Bugs in code, just some Random Unexpected FEATURES ====
┌─────────────────────────────────────────────────────────────┐
│┌───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┬───┐│
││Esc│!1 │@2 │#3 │$4 │%5 │^6 │&7 │*8 │(9 │)0 │_- │+= │|\ │`~ ││
│├───┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴───┤│
││ Tab │ Q │ W │ E │ R │ T │ Y │ U │ I │ O │ P │{[ │}] │ BS  ││
│├─────┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴┬──┴─────┤│
││ Ctrl │ A │ S │ D │ F │ G │ H │ J │ K │ L │: ;│" '│ Enter  ││
│├──────┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴─┬─┴────┬───┤│
││ Shift  │ Z │ X │ C │ V │ B │ N │ M │< ,│> .│? /│Shift │Fn ││
│└─────┬──┴┬──┴──┬┴───┴───┴───┴───┴───┴──┬┴───┴┬──┴┬─────┴───┘│
│      │Fn │ Alt │         Space         │ Alt │Win│   HHKB   │
│      └───┴─────┴───────────────────────┴─────┴───┘          │
└─────────────────────────────────────────────────────────────┘

参照 OpenLlama 改写训练文件。

Authors: pankeyu
Date: 2023/05/25

Reference Authors: s-JoL(sl12160010@gmail.com)
Reference repo: Open-Llama
"""
import os
import time
import shutil

import torch
import torch.nn.functional as F
from torch.optim.lr_scheduler import CosineAnnealingLR

from tqdm import tqdm
from torchinfo import summary
from deepspeed.ops.adam import FusedAdam

import matplotlib.pyplot as plt

from solver.evaluator import *
from iTrainingLogger import iSummaryWriter

class RewardModelTrainer(object):

    def __init__(
            self,
            config,
            raw_model,
            train_loader,
            tokenizer,
            accelerator
        ):
        self.config = config
        self.raw_model = raw_model
        self.train_loader = train_loader
        self.tokenizer = tokenizer
        self.accelerator = accelerator
        self.delta_scores = []
        self.train_and_eval = config["train"].get("train_and_eval", False)
        self.gradient_accumulation_steps = config["train"].get(
            "gradient_accumulation_steps", 1
        )
        self.log_interval = (
            self.config["log_interval"] * accelerator.gradient_accumulation_steps
        )
        self.eval_interval = (
            self.config["eval_interval"] * accelerator.gradient_accumulation_steps
        )
        self.save_interval = (
            self.config["save_interval"] * accelerator.gradient_accumulation_steps
        )
        self.work_dir = self.config["work_dir"]
        if accelerator.is_main_process:
            self.writer = iSummaryWriter(
                config['train'].get('img_log_dir', 'log'),
                config['train'].get('img_log_name', 'performance')
            )

    def get_model_info(self):
        with torch.no_grad():
            summary(
                self.raw_model.cuda(),
                input_data=torch.ones(1, 64, dtype=torch.int64).cuda(),
            )

    def get_optimizer(self):
        no_decay = ["bias", "LayerNorm.weight", "layernorm.weight"]

        if self.config["train"].get("use_lora", False):
            optimizer_grouped_parameters = self.raw_model.parameters()
        else:
            optimizer_grouped_parameters = [
                {
                    "params": [
                        p
                        for n, p in self.raw_model.named_parameters()
                        if not any(nd in n for nd in no_decay)
                        and p.requires_grad
                    ],
                    "weight_decay": self.config["train"]["weight_decay"],
                },
                {
                    "params": [
                        p
                        for n, p in self.raw_model.named_parameters()
                        if any(nd in n for nd in no_decay)
                        and p.requires_grad
                    ],
                    "weight_decay": 0.0,
                },
            ]

        self.optim = FusedAdam(
            optimizer_grouped_parameters,
            lr=self.config["train"]["lr"],
            betas=(0.9, 0.95),
        )

    def get_lr_scheduler(self):
        self.scheduler = CosineAnnealingLR(
            self.optim,
            T_max=len(self.train_loader) * self.config["train"]["num_training_epochs"],
            eta_min=self.config["train"]["min_lr"]
        )

    def prepare(self):
        """
        用于 accelerate 更新。
        """
        (
            train_loader,
            self.model,
            self.optim,
            self.scheduler,
        ) = self.accelerator.prepare(
            self.train_loader,
            self.raw_model,
            self.optim,
            self.scheduler
        )

        self.optim.zero_grad()
        self.global_step = 0
        self.train_loader = train_loader
        self.accelerator.wait_for_everyone()

    def save_model(self):
        """
        存储模型，当存储路径的模型数超过阈值时，删除过期的模型。
        """
        if self.accelerator.is_main_process:
            if not os.path.exists(self.work_dir):
                os.makedirs(self.work_dir)

            while len(os.listdir(self.work_dir)) >= self.config['train'].get('save_total_limit', 100):
                all_checkpoints_files = sorted(                                                 # 按照文件修改时间排序
                    os.listdir(self.work_dir),
                    key=lambda x: os.path.getmtime(os.path.join(self.work_dir, x))
                )
                outdated_model = os.path.join(self.work_dir, all_checkpoints_files[0])
                shutil.rmtree(outdated_model)
                self.accelerator.print(f'[-] Deleted Outdated Model: {outdated_model}.')

        start = time.time()
        deep_speed_stage = os.environ.get('ACCELERATE_DEEPSPEED_ZERO_STAGE', '')
        stage3_save_16bit_model = os.environ.get('ACCELERATE_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL', 'false')

        if stage3_save_16bit_model in ['true', 'True']:
            stage3_save_16bit_model = True
        elif stage3_save_16bit_model in ['false', 'False']:
            stage3_save_16bit_model = False
        else:
            raise ValueError(f'`stage3_save_16bit_model` in accelerate config should be a bool str, while reveived: {stage3_save_16bit_model}.')

        if deep_speed_stage == '3':                                        # Stage3 需要走单独存储的方式
            if stage3_save_16bit_model:
                unwrapped_model = self.accelerator.unwrap_model(self.model)
                cur_save_dir = os.path.join(self.work_dir, f'checkpoint-{self.data_step}')
                unwrapped_model.save_pretrained(
                    cur_save_dir,
                    is_main_process=self.accelerator.is_main_process,
                    save_function=self.accelerator.save,
                    state_dict=self.accelerator.get_state_dict(self.model)
                )
                self.accelerator.print(f'[+] Stage {deep_speed_stage}: Checkpoint has saved at {cur_save_dir}.')
            else:
                self.model.save_checkpoint(self.work_dir, f'checkpoint-{self.data_step}')
                self.accelerator.print(f'[+] Stage 3: Checkpoint has saved at {os.path.join(self.work_dir, f"checkpoint-{self.data_step}")}.')
            self.accelerator.print(f'Save ckpt used: {round(time.time() - start, 2)}s.')
        else:
            if self.accelerator.is_main_process:
                unwrapped_model = self.accelerator.unwrap_model(self.model)
                cur_save_dir = os.path.join(self.work_dir, f'checkpoint-{self.data_step}')
                unwrapped_model.save_pretrained(
                    cur_save_dir,
                    is_main_process=self.accelerator.is_main_process,
                    save_function=self.accelerator.save,
                    state_dict=self.accelerator.get_state_dict(self.model)
                )
                self.accelerator.print(f'[+] Stage {deep_speed_stage}: Checkpoint has saved at {cur_save_dir}.')
                self.accelerator.print(f'Save ckpt used: {round(time.time() - start, 2)}s.')

    def train_step(self, batch):
        """
        单步更新。
        """
        scores = self.model(**batch)[0]                 # (2 * batch, 1)
        scores = scores.reshape(-1, 2).diff()           # (batch, 1), tensor.diff 是后项 - 前项, 因此前面数据准备时是 (refject, select)
        loss = -F.logsigmoid(scores).mean()             # e.g. torch.tensor([[1, 2, 3], [3, 4, 5]]).diff() -> tensor([[2, 2, 2]])
        losses = {"total_loss": loss}
        self.accelerator.backward(loss)
        self.optim.step()
        self.scheduler.step()
        self.optim.zero_grad()
        return losses

    def train(self):
        """
        训练loop。
        """
        self.get_optimizer()
        self.get_lr_scheduler()
        self.prepare()
        self.start_time = time.time()
        self.data_step = 0

        loss_dict_list = {}
        with tqdm(total=self.config["train"]["num_training_epochs"] * len(self.train_loader)) as pbar:
            for ipeoch in range(self.config["train"]["num_training_epochs"]):
                for batch in self.train_loader:

                    if (
                        self.data_step % self.eval_interval == 0
                        and self.train_and_eval
                    ):
                        self.eval()

                    if (
                        self.data_step % self.save_interval == 0
                    ):
                        self.save_model()

                    for k, v in batch.items():
                        batch[k] = v.to(
                            self.accelerator.device,
                            non_blocking=True
                        )

                    self.model.train()
                    with self.accelerator.accumulate(self.model):
                        losses = self.train_step(batch)

                        if self.accelerator.is_main_process:
                            for k, v in losses.items():
                                if k not in loss_dict_list:
                                    loss_dict_list[k] = []
                                loss_dict_list[k].append(v.cpu().item())

                        if self.accelerator.sync_gradients:
                            self.global_step += 1
                            self.accelerator.clip_grad_norm_(
                                self.model.parameters(),
                                1.0
                            )

                    if (
                        self.data_step % self.log_interval == 0
                        and self.accelerator.is_main_process
                    ):
                        self.local_log(loss_dict_list)
                        loss_dict_list = {}

                    self.data_step += 1
                    if self.accelerator.is_main_process:
                        pbar.update(1)

                self.accelerator.print(f'Start {ipeoch} -> {ipeoch + 1} epoch...')

    def local_log(
        self,
        loss_dict_list
    ):
        """
        log to local file.

        Args:
            loss_dict_list (_type_): {
                'total_loss': [7.878, 6.79, 6.92, ...]
            }
        """
        cost_time = time.time() - self.start_time
        self.start_time = time.time()
        tokens = (
            self.config["data"]["batch_size"]
            * self.log_interval
            * self.config["data"]["seq_length"]
        )

        token_per_second = 0 if self.data_step == 0 else tokens / cost_time
        self.writer.add_scalar(
            f"Training/Token per second per gpu",
            token_per_second,
            self.data_step
        )
        self.accelerator.print(f'token per second per gpu: {token_per_second}')

        for k, values in loss_dict_list.items():
            self.writer.add_scalar(
                "Training/{}".format(k),
                sum(values) / len(values),
                self.data_step
            )
            self.accelerator.print(f'{k}: {sum(values) / len(values)}')

        current_lr = self.optim.param_groups[0]["lr"]

        self.writer.add_scalar(
            "Training/LR",
            current_lr,
            self.data_step
        )
        self.accelerator.print(f"lr: {current_lr}")

        current_tokens = (
            self.config["data"]["batch_size"]
            * self.data_step
            * self.config["data"]["seq_length"]
        )
        current_tokens = current_tokens / 1e9               # 统计已训练的token数，转换为B单位

        self.writer.add_scalar(
            "Training/Trained tokens per gpu(B)",
            current_tokens,
            self.data_step
        )
        self.accelerator.print(f"trained tokens per gpu(B): {current_tokens}")

        if self.optim.scaler is not None:
            self.writer.add_scalar(
                "Training/Loss Scale",
                self.optim.scaler.get_scale(),
                self.data_step
            )
            self.accelerator.print(f"loss scale: {self.optim.scaler.get_scale()}")

        self.writer.record()

    def draw_score_delta(self):
        """
        绘制 score delta 的分布。
        """
        plt.xlabel('Score Delta')
        plt.ylabel('Count(s)')
        for i, delta_score in enumerate(self.delta_scores):
            lablel = 'First Score Delta' if not i else 'Current Score Delta'
            plt.hist(
                delta_score,
                50,
                alpha=0.5,
                label=lablel
            )
        plt.legend()

    def eval(self):
        """
        评测函数。
        """
        self.accelerator.wait_for_everyone()                   # stage 3 needs all model params on all devices

        if self.accelerator.is_main_process:
            self.accelerator.print(f'Start to eval reward model...')

        self.model.eval()
        metrics_dict = eval_reward_model_pair_acc(
            self.model,
            self.tokenizer,
            self.accelerator.device,
            self.config['eval']['test_reward_model_acc_files'],
            self.config["data"]["seq_length"]
        )

        if self.accelerator.is_main_process:
            delta_scores = metrics_dict.pop('delta_scores')
            for k, v in metrics_dict.items():
                self.writer.add_scalar(
                    f"eval/RM/{k}",
                    v,
                    self.data_step
                )
            self.accelerator.print(f'RewardModel Eval: {metrics_dict}')
            self.writer.record()

            if self.config['eval'].get('save_delta_scores', False):
                save_path = self.config['eval'].get('delta_scores_save_path', 'delta_scores.json')
                if not os.path.exists(save_path):
                    os.makedirs(save_path)
                    json.dump({}, open(os.path.join(save_path, 'delta_scores.json'), 'w', encoding='utf8'))
                delta_score_dict = json.load(open(os.path.join(save_path, 'delta_scores.json'), 'r', encoding='utf8'))
                delta_score_dict[self.data_step] = delta_scores
                json.dump(delta_score_dict, open(os.path.join(save_path, 'delta_scores.json'), 'w', encoding='utf8'))

        self.accelerator.wait_for_everyone()
        self.model.train()


from django.urls import path, include
from . import views

urlpatterns = [
    path('',views.homePage,name='home'),
    path('accounts/',include('allauth.urls')),
    path('dashboard',views.dashbo,name='dashboard'),
    path('profile/<str:pk>', views.profile, name='users-profile'),
    path('update-user',views.update,name='update-user'),

]


from flask import render_template, flash, redirect, url_for, request, jsonify
from flask_login import current_user, login_user, logout_user, login_required
from app import app, db, photos
from app.forms import LoginForm, RegistrationForm, NewCharacterForm, CharacterSearchForm, DiceRoll, NewPartyForm
from app.models import User, Character, Sub_Race, Sub_Class, Results
from werkzeug.urls import url_parse
from werkzeug.utils import secure_filename
from pathlib import Path, PureWindowsPath, PurePath
import json
import os

@app.route('/', methods=['GET', 'POST'])
@app.route('/index', methods=['GET', 'POST'])
@login_required
def index():
	characters = Character.query.all()
	return render_template('index.html', title='Home Page', characters=characters)

@app.route('/login', methods=['GET', 'POST'])
def login():
    if current_user.is_authenticated:
        return redirect(url_for('index'))
    form = LoginForm()
    if form.validate_on_submit():
        user = User.query.filter_by(username=form.username.data).first()
        if user is None or not user.check_password(form.password.data):
            flash('Invalid username or password')
            return redirect(url_for('login'))
        login_user(user, remember=form.remember_me.data)
        next_page = request.args.get('next')
        if not next_page or url_parse(next_page).netloc != '':
            next_page = url_for('index')
        return redirect(next_page)
    return render_template('login.html', title='Sign In', form=form)

@app.route('/logout')
def logout():
	logout_user()
	return redirect(url_for('index'))

@app.route('/register', methods=['GET', 'POST'])
def register():
    if current_user.is_authenticated:
        return redirect(url_for('index'))
    form = RegistrationForm()
    if form.validate_on_submit():
        user = User(username=form.username.data, first_name=form.first_name.data)
        user.set_password(form.password.data)
        db.session.add(user)
        db.session.commit()
        flash('Congratulations, you are now a registered user!')
        return redirect(url_for('login'))
    return render_template('register.html', title='Register', form=form)

@app.route('/user/<username>')
@login_required
def user(username):
    user = User.query.filter_by(username=username).first_or_404()
    characters = current_user.get_characters()
    return render_template('user.html', user=user, characters=characters, title='Profile')

@app.route('/user/<username>/create_new_char', methods=['GET', 'POST'])
@login_required
def create_new_char(username):
	user = User.query.filter_by(username=username).first_or_404()
	form = NewCharacterForm()
	form.sub_class.choices = [(sub_class.name, sub_class.name) for sub_class in Sub_Class.query.filter_by(char_class='Barbarian').all()]
	form.sub_race.choices = [(sub_race.name, sub_race.name) for sub_race in Sub_Race.query.filter_by(race='Aarakocra').all()]
	if request.method == 'POST':
		character = Character(name=form.name.data, race=form.race.data,\
		 char_class=form.char_class.data, sub_class=form.sub_class.data, sub_race=form.sub_race.data, level=form.level.data, hp=form.hp.data, user_id=user.id)
		db.session.add(character)
		db.session.commit()
		flash('Your character has been created!')
		return redirect(url_for('user', username=current_user.username))
	return render_template('create_new_char.html', user=user, form=form, title='Create New Character')

@app.route('/user/<username>/edit_character/<character>', methods=['GET', 'POST'])
@login_required
def edit_character(username, character):
	user = User.query.filter_by(username=username).first_or_404()
	char_id = user.user_id
	char = Character.query.filter_by(user_id=char_id, name=character).first_or_404()
	form = NewCharacterForm(obj=char)
	current_char_race = char.race
	current_char_class = char.char_class
	form.sub_race.choices = [(sub_race.name, sub_race.name) for sub_race in Sub_Race.query.filter_by(race=current_char_race).all()]
	form.sub_class.choices = [(sub_class.name, sub_class.name) for sub_class in Sub_Class.query.filter_by(char_class=current_char_class).all()]
	if request.method == 'POST':
		form.populate_obj(char)
		db.session.add(char)
		db.session.commit()
		flash('Your character has been saved!')
		return redirect(url_for('user', username=current_user.username))
	return render_template('edit_character.html', user=user, form=form, title='Edit Character')

@app.route('/user/<username>/<character>', methods=['GET', 'POST'])
@login_required
def display_character(username, character):
	if character == None:
		return redirect(url_for('index'))
	user = User.query.filter_by(username=username).first_or_404()
	char_id = user.id
	char = Character.query.filter_by(user_id=char_id, name=character).first_or_404()
	return render_template('display_char.html', username=user, character=char, title=str(character))

@app.route('/search', methods=['GET', 'POST'])
@login_required
def search():
	search = CharacterSearchForm(request.form)
	if request.method == 'POST':
		return search_results(search)
	return render_template('search.html', form=search, title="Search")

@app.route('/results', methods=['GET', 'POST'])
@login_required
def search_results(search):
	results = []
	first_string = search.data['search']
	second_string = first_string.lower()
	search_string = second_string.capitalize()
	if search_string:
		if search.data['select'] == 'name':
			qry = Character.query.filter(Character.name.contains(search_string))
			results = qry.all()
		elif search.data['select'] == 'char_class':
			qry = Character.query.filter(Character.char_class.contains(search_string))
			results = qry.all()
		elif search.data['select'] == 'race':
			qry = Character.query.filter(Character.race.contains(search_string))
			results = qry.all()
		elif search.data['select'] == 'level':
			qry = Character.query.filter(Character.level.contains(search_string))
			results = qry.all()
		# elif search.data['select'] == 'party':
		# 	qry = Character.query.filter(Character.party.contains(search_string))
		# 	results = qry.all()
		else:
			qry = db.session.query(Character)
			results = qry.all()
	else:
		qry = db.session.query(Character)
		results = qry.all()

	if not results:
		flash('No results found!')
		return redirect(url_for('search'))

	else:
		table = Results(results)
		table.border = True
		return render_template('results.html', results=results, table=table, title="Search Results")

@app.route('/dice_roller', methods=['GET', 'POST'])
@login_required
def dice_roller():
	form = DiceRoll(request.form)
	num = 0
	return render_template('dice_roller.html', form=form, num=num)

@app.route('/dice_update', methods=['POST'])
def dice_update():
	num = request.form['result']
	return jsonify({'result': num})

def allowed_file(filename):
	return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

# create new party
@app.route('/create_new_party', methods=['GET', 'POST'])
@login_required
def create_new_party():
	user = User.query.filter_by(username=current_user.username).first_or_404()
	form = NewPartyForm()
	form.party_members.choices = [(character.name, character.name) for character in Character.query.all()]

	if request.method == 'POST':
		current_party_list = party_list()
		new_party = Party(party_name=form.party_name.data, party_leader=user.id)
		for char in current_party_list:
			char.party_id = new_party.party_id
		db.session.add(new_party)
		db.session.commit()
		flash('Your party has been created!')
		return redirect(url_for('user', username=current_user.username))
	return render_template('create_new_party.html', form=form, user=user)

@app.route('/party_list')
@login_required
def party_list():
	current_party_list = json.loads(request.args.get('party_list'))
	return jsonify(result=current_party_list)

@app.route('/upload/<username>', methods=['GET', 'POST'])
@login_required
def upload(username):
	user = User.query.filter_by(username=username).first_or_404()

	if request.method == 'POST':

		if 'photo' not in request.files:
			flash('No file part')
			return redirect(url_for('user', username=username))

		file = request.files['photo']

		if file.filename == '':
			flash("No selected file")
			return redirect(url_for('user', username=username))

		if file and allowed_file(file.filename):

			# TODO
			# language to delete old photos using OS neutral path
			# if user.profile_photo:
			# 	delete_old = os.path.normpath(user.profile_photo)
			# 	flash(str(delete_old))
			# 	os.remove(user.profile_photo)

			# save requested img to server
			filename = Path(secure_filename(file.filename)).as_posix()
			filename_2 = photos.save(request.files['photo'])

			# establish OS neutral file path to save to DB
			pure_path = Path(PureWindowsPath('\\static\\img\\'))
			user.profile_photo = os.path.join(pure_path.as_posix(), filename)

			#commit filepath string to db
			db.session.commit()

			flash("New Avatar Uploaded Sucessfully!")
			return redirect(url_for('user', username=username))
	return render_template('upload.html', username=username)

# dynamic SelectField for sub_races
@app.route('/sub_race/<race>')
@login_required
def sub_race(race):
	sub_races = Sub_Race.query.filter_by(race=race).all()
	sub_race_array = []
	for sub_race in sub_races:
		sub_race_obj = {}
		sub_race_obj['name'] = sub_race.name
		sub_race_array.append(sub_race_obj)
	return jsonify({'sub_races' : sub_race_array})

# dynamic SelectField for sub_classes
@app.route('/sub_class/<char_class>')
@login_required
def sub_class(char_class):
	sub_classes = Sub_Class.query.filter_by(char_class=char_class).all()
	sub_class_array = []
	for sub_class in sub_classes:
		sub_class_obj = {}
		sub_class_obj['name'] = sub_class.name
		sub_class_array.append(sub_class_obj)
	return jsonify({'sub_classes' : sub_class_array})

@app.route('/delete_confirm/<character>', methods=['GET', 'POST'])
@login_required
def delete_confirm(character):
	char = Character.query.filter_by(name=character).first_or_404()
	char_name = char.name

	if request.method == 'POST':
		db.session.delete(char)
		db.session.commit()
		flash("Character Successfully Deleted!")
		return redirect(url_for('user', username=current_user.username))

	return render_template('delete_character.html', title="Delete Character", character=char_name)


import json
import csv

with open('entailmentbank/data/chains_train.json','r') as f:
    true_train=json.load(f)

with open('entailmentbank/outputs/pre_train_top50.json','r') as f:
    pre_train=json.load(f)

with open("entailmentbank/data/worldtree_corpus_sentences_extended.json", 'r') as f:
    knowledge_corpus = json.load(f)

with open("entailmentbank/data/hypotheses_train.json", 'r') as f:
    hypotheses_train = json.load(f)

output=open('entailmentbank/train/rerank_train_2th.csv','w')

for i,j in pre_train.items():
    true_k=true_train[i] #true list k
    for k in j:
        if k in true_k:

            print(hypotheses_train[i] + "\t" + knowledge_corpus[k] + "\t" + str(1), file=output)
        else:
            print(hypotheses_train[i] + "\t" + knowledge_corpus[k] + "\t" + str(0), file=output)

    # print(i,j)
    # break


import unittest
from disjoint_set import DisjointSet
"""
Time complexity: O(E * LogE)
Space complexity: O(V)
"""
def kruskals(n, edges):
  ds = DisjointSet()
  edge_count = 0
  total_weight = 0
  for w, a, b in sorted(edges):
    ds.add(a)
    ds.add(b)
    if ds.union(a, b):
      edge_count += 1
      total_weight += w
      if edge_count == n-1:
        break
  return total_weight

class TestKruskals(unittest.TestCase):
  def test_kruskals(self):
    graph = {
      'A': [(2, 'B'), (3, 'C')],
      'B': [(2, 'A'), (1, 'C'), (1, 'D'), (4, 'E')],
      'C': [(3, 'A'), (1, 'B'), (5, 'F')],
      'D': [(1, 'B'), (1, 'E')],
      'E': [(4, 'B'), (1, 'D'), (1, 'F')],
      'F': [(5, 'C'), (1, 'E'), (1, 'G')],
      'G': [(1, 'F')],
    }

    edges = []
    for a in graph:
      for w, b in graph[a]:
        edges.append((w, a, b))
    self.assertEqual(kruskals(len(graph), edges), 7)
if __name__ == "__main__":
  unittest.main()


from stellargraph.layer import (
    DeepGraphInfomax,
    GCN,
    APPNP,
    GAT,
    PPNP,
    GraphSAGE,
    DirectedGraphSAGE,
)

from stellargraph.mapper import (
    FullBatchNodeGenerator,
    CorruptedGenerator,
    GraphSAGENodeGenerator,
    DirectedGraphSAGENodeGenerator,
    FullBatchLinkGenerator,
)
from ..test_utils.graphs import example_graph_random
import tensorflow as tf
import pytest
import numpy as np

@pytest.mark.parametrize("model_type", [GCN, APPNP, GAT, PPNP])
@pytest.mark.parametrize("sparse", [False, True])
def test_dgi(model_type, sparse):

    if sparse and model_type is PPNP:
        pytest.skip("PPNP doesn't support sparse=True")

    G = example_graph_random()
    emb_dim = 16

    generator = FullBatchNodeGenerator(G, sparse=sparse)
    corrupted_generator = CorruptedGenerator(generator)
    gen = corrupted_generator.flow(G.nodes())

    base_model = model_type(
        generator=generator, activations=["relu"], layer_sizes=[emb_dim]
    )
    infomax = DeepGraphInfomax(base_model)

    model = tf.keras.Model(*infomax.in_out_tensors())
    model.compile(loss=tf.nn.sigmoid_cross_entropy_with_logits, optimizer="Adam")
    model.fit(gen)

    emb_model = tf.keras.Model(*infomax.embedding_model())
    embeddings = emb_model.predict(generator.flow(G.nodes()))

    assert embeddings.shape == (len(G.nodes()), emb_dim)

def test_dgi_stateful():
    G = example_graph_random()
    emb_dim = 16

    generator = FullBatchNodeGenerator(G)
    corrupted_generator = CorruptedGenerator(generator)
    gen = corrupted_generator.flow(G.nodes())

    infomax = DeepGraphInfomax(
        GCN(generator=generator, activations=["relu"], layer_sizes=[emb_dim])
    )

    model_1 = tf.keras.Model(*infomax.in_out_tensors())
    model_2 = tf.keras.Model(*infomax.in_out_tensors())

    # check embeddings are equal before training
    embeddings_1 = tf.keras.Model(*infomax.embedding_model()).predict(
        generator.flow(G.nodes())
    )
    embeddings_2 = tf.keras.Model(*infomax.embedding_model()).predict(
        generator.flow(G.nodes())
    )

    assert np.array_equal(embeddings_1, embeddings_2)

    model_1.compile(loss=tf.nn.sigmoid_cross_entropy_with_logits, optimizer="Adam")
    model_1.fit(gen)

    # check embeddings are still equal after training one model
    embeddings_1 = tf.keras.Model(*infomax.embedding_model()).predict(
        generator.flow(G.nodes())
    )
    embeddings_2 = tf.keras.Model(*infomax.embedding_model()).predict(
        generator.flow(G.nodes())
    )

    assert np.array_equal(embeddings_1, embeddings_2)

    model_2.compile(loss=tf.nn.sigmoid_cross_entropy_with_logits, optimizer="Adam")
    model_2.fit(gen)

    # check embeddings are still equal after training both models
    embeddings_1 = tf.keras.Model(*infomax.embedding_model()).predict(
        generator.flow(G.nodes())
    )
    embeddings_2 = tf.keras.Model(*infomax.embedding_model()).predict(
        generator.flow(G.nodes())
    )

    assert np.array_equal(embeddings_1, embeddings_2)

@pytest.mark.parametrize("model_type", [GCN, APPNP, GAT])
def test_dgi_link_model(model_type):
    G = example_graph_random()
    emb_dim = 16

    generator = FullBatchLinkGenerator(G)

    with pytest.warns(
        UserWarning,
        match=r"base_model: expected a node model .* found a link model \(multiplicity = 2\)",
    ):
        infomax = DeepGraphInfomax(
            model_type(generator=generator, activations=["relu"], layer_sizes=[emb_dim])
        )

    # build should work
    _ = infomax.in_out_tensors()

@pytest.mark.parametrize("is_directed", [False, True])
def test_dgi_graphsage(is_directed):

    G = example_graph_random(is_directed=is_directed)

    if is_directed:
        generator = DirectedGraphSAGENodeGenerator(
            G, batch_size=5, in_samples=[2, 3], out_samples=[4, 1]
        )
        base_model = DirectedGraphSAGE(generator=generator, layer_sizes=[4, 4])
    else:
        generator = GraphSAGENodeGenerator(G, batch_size=5, num_samples=[2, 3])
        base_model = GraphSAGE(generator=generator, layer_sizes=[4, 4])

    corrupted_generator = CorruptedGenerator(generator)
    gen = corrupted_generator.flow(G.nodes())

    infomax = DeepGraphInfomax(base_model)

    model = tf.keras.Model(*infomax.build())
    model.compile(loss=tf.nn.sigmoid_cross_entropy_with_logits, optimizer="Adam")
    model.fit(gen)

    emb_model = tf.keras.Model(*infomax.embedding_model())
    embeddings = emb_model.predict(generator.flow(G.nodes()))

    assert embeddings.shape == (len(G.nodes()), 4)