import torch

# set printing options
torch.set_printoptions(sci_mode=False)

# define the a matrix-elementwise metric, the mean square error 
def mse(y_hat, y_true): return ((y_hat-y_true)**2).mean()

def improve(U, diag_sqrt, V, M, learning_rate, momentum, losses):

  # the given matrix
  y_true = M

  # estimate the original matrix M
  diag = diag_sqrt * diag_sqrt 
  # this enforces that the diagonal values are >=0
  
  S = torch.diag(input=diag)

  y_hat =  U @ S @ V.T
  
  # loss = "degree of being lost"
  loss1 = mse(y_hat, y_true) # 

  # further constraints
  # 2) orthonomality U.T @ U = I_U
  I_U = torch.eye(U.shape[1])
  loss2 = mse(U.T @ U, I_U)

  # 3) orthonomality U.T @ U = I_U
  I_V = torch.eye(V.shape[0])
  loss3 = mse(V @ V.T, I_V)

  # add up the loss with some scaling factors 
  # ... to improve convergence
  loss = loss1 + loss2/100 + loss3/100


  losses.append(loss.detach().numpy())


  # calculate loss
  loss.backward()

  with torch.no_grad():
    # displace by learning_rate * derivatives
    
    # (i)

    U -= learning_rate * U.grad
    diag_sqrt -= learning_rate * diag_sqrt.grad
    V -= learning_rate * V.grad

    # remark: instead of (i) we could alternatively use sub_ and write (ii),
    #         this is computationally more efficient but harder to read.
    # (ii)
    # approximate_inverse.sub_(learning_rate * approximate_inverse.grad)
    

    # momentum
    # (iii)
    U.grad *= momentum # reduce speed due to friction
    diag_sqrt.grad *= momentum # reduce speed due to friction
    V.grad *= momentum # reduce speed due to friction

    # remark for the case of momentum == 0 we could alternatively write
    # (iv)
    # approximate_inverse.grad.zero_()


  return None

def calculate_svd(M, learning_rate=0.9, momentum=0.9):

  torch.manual_seed(314)
  # Step 0) Guess
  U = torch.rand(size=M.shape, requires_grad=True)
  diag_sqrt = torch.rand(size=[M.shape[1]], requires_grad=True)
  V = torch.rand(size=(M.shape[1], M.shape[1]), requires_grad=True)
  losses = []
  

  for t in range(20000): 
    # Step 1) Improve
    improve(U, diag_sqrt, V,  M, learning_rate, momentum, losses)
  

  diag = diag_sqrt*diag_sqrt

  # sort diagonal values in descending order
  diag_sorted, order = diag.sort(descending=True)
  diag = diag[order]
  U = U[:,order]
  V = V[:,order]

  return U, diag, V, losses

torch.manual_seed(314)
M1 = torch.rand(size=(5,4)) # torch.diag(torch.tensor([1.,2.,3.,4.]))

U1, diag1, V1, losses_1 = calculate_svd(M1, learning_rate=0.9, momentum=0.9)

S1 = torch.diag(input=diag1)
U1 @S1 @V1.T - M1, U1.T @ U1, V1 @ V1.T

U, S, V = torch.svd(M1)

U , '',U1

S, '',  torch.diag(S1)

V, '' ,V1

import pandas as pd
import numpy as np


import plotly.express as px
df = pd.DataFrame(
    {
     'losses_1': losses_1,
     })
fig = px.line(df, y='losses_1', log_y=True)
fig.update_traces(name='Random Matrix 1, Momentum=0.9', showlegend = True)

fig.update_layout(title=\
                  '<b>How fast does the gradient algorithm converge?</b>',
                  xaxis_title='Iterations',
                  yaxis_title='Losses'
                  )
fig

torch.manual_seed(315)

M2 = torch.rand(size=(4,4)) # torch.diag(torch.tensor([1.,2.,3.,4.]))

_, _, _, losses_2 = calculate_svd(M2, learning_rate=0.9, momentum=0.9)

torch.manual_seed(316)

M3 = torch.rand(size=(7,4)) # torch.diag(torch.tensor([1.,2.,3.,4.]))

_, _, _, losses_3 = calculate_svd(M3, learning_rate=0.9, momentum=0.9)

torch.manual_seed(317)

M4 = torch.rand(size=(3,4)).T # torch.diag(torch.tensor([1.,2.,3.,4.]))

_, _, _, losses_4 = calculate_svd(M4, learning_rate=0.9, momentum=0.9)

#####################################################################
torch.manual_seed(317)

M5 = torch.rand(size=(4,3)) # torch.diag(torch.tensor([1.,2.,3.,4.]))

_, _, _, losses_5 = calculate_svd(M5, learning_rate=0.9, momentum=0.9)

import pandas as pd
import numpy as np


import plotly.express as px
df = pd.DataFrame(
    {
     'losses_1': losses_1,
     'losses_2' : losses_2,
     'losses_3':losses_3,
     'losses_4': losses_4,
     'losses_5':losses_5
     })
fig = px.line(df, y='losses_1', log_y=True)
fig.update_traces(name='Random Matrix 1, Momentum=0.9', showlegend = True)
fig.add_scatter(y=df['losses_2'], name='Random Matrix 2,  Momentum=0.9')
fig.add_scatter(y=df['losses_3'], name='Random Matrix 3,  Momentum=0.9')
fig.add_scatter(y=df['losses_4'], name='Random Matrix 4,  Momentum=0.9')
fig.add_scatter(y=df['losses_5'], name='Random Matrix 5,  Momentum=0.9')

fig.update_layout(title=\
                  '<b>How fast does the gradient algorithm converge?</b>' + \
                  '<br>Answer:It depends',
                  xaxis_title='Iterations',
                  yaxis_title='Losses'
                  )
fig