"""Analyze condition number of the network."""

import numpy as np
import matplotlib.pyplot as plt

# import model

def _get_sparse_mask(nx, ny, non, complex=False, nOR=50):
    """Generate a binary mask.

    The mask will be of size (nx, ny)
    For all the nx connections to each 1 of the ny units, only non connections are 1.

    Args:
        nx: int
        ny: int
        non: int, must not be larger than nx

    Return:
        mask: numpy array (nx, ny)
    """
    mask = np.zeros((nx, ny))

    if not complex:
        mask[:non] = 1
        for i in range(ny):
            np.random.shuffle(mask[:, i])  # shuffling in-place
    return mask.astype(np.float32)

def _get_cond(q, n_orn, n_pn, n_kc, n_kc_claw):
    M = np.random.rand(n_orn, n_pn)
    M_new = M * (1-q) + np.eye(n_orn) * q
    
    # J = np.random.rand(N_PN, N_KC) / np.sqrt(N_PN + N_KC)
    # J = np.random.randn(N_PN, N_KC) / np.sqrt(N_PN + N_KC)
    J = np.random.rand(n_pn, n_kc)
    mask = _get_sparse_mask(n_pn, n_kc, n_kc_claw) / n_kc_claw
    J = J * mask

    K = np.dot(M_new, J)
    
    # cond = np.linalg.cond(K)

    cond = np.linalg.norm(np.linalg.pinv(K)) * np.linalg.norm(K)
    return cond


def get_logcond(q=1, n_orn=50, n_pn=50, n_kc=2500, n_kc_claw=7, n_rep=10):
    conds = [_get_cond(q, n_orn, n_pn, n_kc, n_kc_claw) for i in range(n_rep)]
    return np.mean(np.log10(conds))



def plot_cond_by_q(n_kc=2500):
    qs = np.linspace(0, 1, 100)
    conds = [get_logcond(q=q, n_kc=n_kc) for q in qs]
    
    plt.figure()
    plt.plot(qs, conds, 'o-')
    plt.title('N_KC: ' + str(n_kc))
    plt.xlabel('fraction diagonal')
    plt.ylabel('log condition number')
    # plt.savefig('figures/condvsfracdiag_nkc'+str(n_kc)+'.pdf', transparent=True)


def plot_cond_by_n_kc():
    n_kcs = np.logspace(1, 4, 10).astype(int)
    conds_q1 = np.array([get_logcond(n_kc=n_kc, q=1) for n_kc in n_kcs])
    
    plt.figure()
    plt.plot(np.log10(n_kcs), conds_q1, 'o-')
    plt.xticks(np.log10(n_kcs), n_kcs)
    plt.xlabel('N_KC')
    
    
    n_kcs = np.logspace(1, 4, 10).astype(int)
    conds_q0 = np.array([get_logcond(n_kc=n_kc, q=0) for n_kc in n_kcs])
    
    plt.figure()
    plt.plot(np.log10(n_kcs), conds_q0, 'o-')
    plt.xticks(np.log10(n_kcs), n_kcs)
    plt.xlabel('N_KC')
    
    plt.figure()
    plt.plot(np.log10(n_kcs), conds_q1 - conds_q0, 'o-')
    plt.xticks(np.log10(n_kcs), n_kcs)
    plt.ylabel('Log decrease in condition number')
    plt.xlabel('N_KC')
    
    
n_kc_claws = np.arange(1, 50)
conds = np.array([get_logcond(n_kc_claw=n) for n in n_kc_claws])

plt.figure()
plt.plot(n_kc_claws, conds, 'o-')
plt.xticks(n_kc_claws)
plt.xlabel('N_KC_claw')
plt.show()

[(91, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (94, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (95, 'matplotlib.pyplot.plot', 'plt.plot', True, 'import matplotlib.pyplot as plt\n'), (96, 'matplotlib.pyplot.xticks', 'plt.xticks', True, 'import matplotlib.pyplot as plt\n'), (97, 'matplotlib.pyplot.xlabel', 'plt.xlabel', True, 'import matplotlib.pyplot as plt\n'), (98, 'matplotlib.pyplot.show', 'plt.show', True, 'import matplotlib.pyplot as plt\n'), (22, 'numpy.zeros', 'np.zeros', True, 'import numpy as np\n'), (31, 'numpy.random.rand', 'np.random.rand', True, 'import numpy as np\n'), (36, 'numpy.random.rand', 'np.random.rand', True, 'import numpy as np\n'), (40, 'numpy.dot', 'np.dot', True, 'import numpy as np\n'), (55, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (58, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (59, 'matplotlib.pyplot.plot', 'plt.plot', True, 'import matplotlib.pyplot as plt\n'), (61, 'matplotlib.pyplot.xlabel', 'plt.xlabel', True, 'import matplotlib.pyplot as plt\n'), (62, 'matplotlib.pyplot.ylabel', 'plt.ylabel', True, 'import matplotlib.pyplot as plt\n'), (70, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (73, 'matplotlib.pyplot.xlabel', 'plt.xlabel', True, 'import matplotlib.pyplot as plt\n'), (79, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (82, 'matplotlib.pyplot.xlabel', 'plt.xlabel', True, 'import matplotlib.pyplot as plt\n'), (84, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (87, 'matplotlib.pyplot.ylabel', 'plt.ylabel', True, 'import matplotlib.pyplot as plt\n'), (88, 'matplotlib.pyplot.xlabel', 'plt.xlabel', True, 'import matplotlib.pyplot as plt\n'), (44, 'numpy.linalg.norm', 'np.linalg.norm', True, 'import numpy as np\n'), (50, 'numpy.log10', 'np.log10', True, 'import numpy as np\n'), (71, 'numpy.log10', 'np.log10', True, 'import numpy as np\n'), (72, 'numpy.log10', 'np.log10', True, 'import numpy as np\n'), (80, 'numpy.log10', 'np.log10', True, 'import numpy as np\n'), (81, 'numpy.log10', 'np.log10', True, 'import numpy as np\n'), (85, 'numpy.log10', 'np.log10', True, 'import numpy as np\n'), (86, 'numpy.log10', 'np.log10', True, 'import numpy as np\n'), (27, 'numpy.random.shuffle', 'np.random.shuffle', True, 'import numpy as np\n'), (32, 'numpy.eye', 'np.eye', True, 'import numpy as np\n'), (44, 'numpy.linalg.pinv', 'np.linalg.pinv', True, 'import numpy as np\n'), (67, 'numpy.logspace', 'np.logspace', True, 'import numpy as np\n'), (76, 'numpy.logspace', 'np.logspace', True, 'import numpy as np\n')]