# Run this cell using Ctrl+Enter (⌘+Enter on Mac).
from testing import exercise, create_empty_matrix
from typing import List

import math, cmath

Matrix = List[List[complex]]

@exercise
def matrix_add(a : Matrix, b : Matrix) -> Matrix:
    # You can get the size of a matrix like this:
    rows = len(a)
    columns = len(a[0])
    
    # You can use the following function to initialize a rows×columns matrix filled with 0s to store your answer
    c = create_empty_matrix(rows, columns)
        
    for i in range(rows):
        for j in range(columns):
            # You can access elements of a matrix like this:
            x = a[i][j]
            y = b[i][j]
            
            # You can modify the elements of a matrix like this:
            c[i][j] = a[i][j] + b[i][j]
    
    return c

@exercise
def scalar_mult(x : complex, a : Matrix) -> Matrix:
    rows = len(a)
    columns = len(a[0])
    
    c = create_empty_matrix(rows, columns)
    
    for i in range(rows):
        for j in range(columns):
            c[i][j] = a[i][j] * x
    
    return c

@exercise
def matrix_mult(a : Matrix, b : Matrix) -> Matrix:
    rows = len(a)        # the number of rows of the left matrix
    common = len(a[0])   # = len(b) - the common dimension of the matrices
    columns = len(b[0])  # the number of columns of the right matrix
    
    ans = create_empty_matrix(rows, columns)
    
    for currentRow in range(rows):
        for currentColumn in range(columns):
            for k in range(common):
                ans[currentRow][currentColumn] += a[currentRow][k] * b[k][currentColumn]
                
    return ans

@exercise
def matrix_inverse(m : Matrix) -> Matrix:
    # Extract each element of the array into a named variable
    a = m[0][0]
    b = m[0][1]
    c = m[1][0]
    d = m[1][1]
    
    # Calculate the determinant
    determinant = (a * d) - (b * c)
    
    # Create the inverse of the matrix following the formula above
    ans = [[d / determinant, -b / determinant], [-c / determinant, a / determinant]]

    return ans

@exercise
def transpose(a : Matrix) -> Matrix:
    rows = len(a)
    columns = len(a[0])
    
    # Note that the resulting matrix dimensions are swapped compared to the original ones
    ans = create_empty_matrix(columns, rows)
    
    for i in range(rows):
        for j in range(columns):
            ans[j][i] = a[i][j]
                
    return ans

@exercise
def conjugate(a : Matrix) -> Matrix:
    rows = len(a)
    columns = len(a[0])
    
    ans = create_empty_matrix(rows, columns)
    
    for i in range(rows):
        for j in range(columns):
            ans[i][j] = complex(a[i][j].real, -a[i][j].imag)
            
    return ans

@exercise
def adjoint(a : Matrix) -> Matrix:
    
    # Call the transpose function with the input matrix a
    transp = transpose(a)
    
    # Call the conjugate function with the transposed matrix as input
    ans = conjugate(transp)
    
    return ans

from pytest import approx

@exercise
def is_matrix_unitary(a : Matrix) -> bool:
    n = len(a)
    
    # Calculate the adjoint matrix
    adjointA = adjoint(a)
    
    # Multiply the adjoint matrix by the input matrix
    multipliedMatrix = matrix_mult(a, adjointA)
    
    # Check whether the multiplication result is (approximately) identity matrix
    for i in range(n):
        for j in range(n):
            # An identity matrix has 1's in all the places where the row index and column index are equal...
            if i == j:
                if multipliedMatrix[i][j] != approx(1): 
                    return False
            # ... and 0's in all the places where the row index and column index are different
            else:
                if multipliedMatrix[i][j] != approx(0):
                    return False
    
    return True

@exercise
def inner_prod(v : Matrix, w : Matrix) -> complex:
    
    # Calculate the adjoint of the v vector
    adjointV = adjoint(v)
    
    # Multiply the adjoint v and w. The result will be a matrix with only one element.
    resultMatrix = matrix_mult(adjointV, w)
    
    # To get the actual complex number, we have to take one element from the multiplication result.
    return resultMatrix[0][0]

@exercise
def normalize(v : Matrix) -> Matrix:
    norm = math.sqrt(inner_prod(v, v).real)
    
    n = len(v)
    ans = create_empty_matrix(n, 1)
    
    # Divide each element of the vector by the norm
    for i in range(n):
        ans[i][0] = v[i][0] / norm
        
    return ans

@exercise
def outer_prod(v : Matrix, w : Matrix) -> Matrix:
    # Calculate adjoint of the W
    adjointW = adjoint(w)
    
    # Multiply V by W adjoint
    return matrix_mult(v, adjointW)

@exercise
def tensor_product(a : Matrix, b : Matrix) -> Matrix:
    aRows = len(a)       # the number of rows for matrix a
    aColumns = len(a[0]) # the number of columns for matrix a
    bRows = len(b)       # the number of rows for matrix b
    bColumns = len(b[0]) # the number of columns for matrix b
    
    ans = create_empty_matrix(aRows * bRows, aColumns *  bColumns)
    
    # Outer pair of loops, iterating trough the elements of the left matrix
    for i in range(aRows):
        for j in range(aColumns):
            # Inner pair of loops, iterating through the elements of the right matrix
            for k in range(bRows):
                for l in range(bColumns):
                    ans[i * bRows + k][j * bColumns + l] = a[i][j] * b[k][l]
    
    return ans

from pytest import approx

@exercise
def find_eigenvalue(a : Matrix, v : Matrix) -> float:
    n = len(v)
    multiplied = matrix_mult(a, v)

    for i in range(n):
        if (v[i][0] != approx(0)):
            return multiplied[i][0] / v[i][0]

@exercise
def find_eigenvector(a : Matrix, x : float) -> Matrix:
    # Check for possible edge cases
    if (a[0][1] == 0):
        if (a[0][0] - x == 0):
            if (a[1][0] == 0):
                return [[1], [0]]
            else:
                return [[(a[1][1] - x) / (-a[1][0])], [1]]
        else:
            return [[0], [1]]
    
    v0 = 1
    v1 = (a[0][0] - x) / (-a[0][1])
    return [[v0], [v1]]