# Important note: you do not have to modify this file for your homework.

import util
import numpy as np


def calc_grad(X, Y, theta):
    """Compute the gradient of the loss with respect to theta."""
    count, _ = X.shape

    probs = 1. / (1 + np.exp(-X.dot(theta)))
    grad = (Y - probs).dot(X)

    return grad


def logistic_regression(X, Y):
    """Train a logistic regression model."""
    theta = np.zeros(X.shape[1])
    learning_rate = 0.1

    i = 0
    while True:
        i += 1
        prev_theta = theta
        theta = theta_step(X, Y, theta, i, learning_rate)
        if np.linalg.norm(prev_theta - theta) < 1e-15:
            print('Converged in %d iterations' % i)
            break
    return

def theta_step(X, Y, theta, i, learning_rate):
    prev_theta = theta
    grad = calc_grad(X, Y, theta)
    theta = theta + learning_rate * grad
    if i % 10000 == 0:
        print('Finished %d iterations' % i)
        print(f'Theta: {theta}')
        print(f'Size of theta: {theta @ theta}')
        print(f'Size of grad: {grad @ grad}')
    return theta

def theta_step_n(X, Y, theta, learning_rate, n):
    for i in range(n):
        theta = theta_step(X, Y, theta, i, learning_rate)
    return theta

def main():
    print('==== Training model on data set A ====')
    Xa, Ya = util.load_csv('ds1_a.csv', add_intercept=True)
    logistic_regression(Xa, Ya)

    print('\n==== Training model on data set B ====')
    Xb, Yb = util.load_csv('ds1_b.csv', add_intercept=True)
    logistic_regression(Xb, Yb)


if __name__ == '__main__':
    main()
