"""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, 13, 91, 29), 'numpy.arange', 'np.arange', ({(91, 23, 91, 24): '1', (91, 26, 91, 28): '50'}, {}), '(1, 50)', True, 'import numpy as np\n'), ((94, 0, 94, 12), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((95, 0, 95, 33), 'matplotlib.pyplot.plot', 'plt.plot', ({(95, 9, 95, 19): 'n_kc_claws', (95, 21, 95, 26): 'conds', (95, 28, 95, 32): '"""o-"""'}, {}), "(n_kc_claws, conds, 'o-')", True, 'import matplotlib.pyplot as plt\n'), ((96, 0, 96, 22), 'matplotlib.pyplot.xticks', 'plt.xticks', ({(96, 11, 96, 21): 'n_kc_claws'}, {}), '(n_kc_claws)', True, 'import matplotlib.pyplot as plt\n'), ((97, 0, 97, 23), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(97, 11, 97, 22): '"""N_KC_claw"""'}, {}), "('N_KC_claw')", True, 'import matplotlib.pyplot as plt\n'), ((98, 0, 98, 10), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((22, 11, 22, 29), 'numpy.zeros', 'np.zeros', ({(22, 20, 22, 28): '(nx, ny)'}, {}), '((nx, ny))', True, 'import numpy as np\n'), ((31, 8, 31, 35), 'numpy.random.rand', 'np.random.rand', ({(31, 23, 31, 28): 'n_orn', (31, 30, 31, 34): 'n_pn'}, {}), '(n_orn, n_pn)', True, 'import numpy as np\n'), ((36, 8, 36, 34), 'numpy.random.rand', 'np.random.rand', ({(36, 23, 36, 27): 'n_pn', (36, 29, 36, 33): 'n_kc'}, {}), '(n_pn, n_kc)', True, 'import numpy as np\n'), ((40, 8, 40, 24), 'numpy.dot', 'np.dot', ({(40, 15, 40, 20): 'M_new', (40, 22, 40, 23): 'J'}, {}), '(M_new, J)', True, 'import numpy as np\n'), ((55, 9, 55, 31), 'numpy.linspace', 'np.linspace', ({(55, 21, 55, 22): '0', (55, 24, 55, 25): '1', (55, 27, 55, 30): '100'}, {}), '(0, 1, 100)', True, 'import numpy as np\n'), ((58, 4, 58, 16), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((59, 4, 59, 29), 'matplotlib.pyplot.plot', 'plt.plot', ({(59, 13, 59, 15): 'qs', (59, 17, 59, 22): 'conds', (59, 24, 59, 28): '"""o-"""'}, {}), "(qs, conds, 'o-')", True, 'import matplotlib.pyplot as plt\n'), ((61, 4, 61, 35), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(61, 15, 61, 34): '"""fraction diagonal"""'}, {}), "('fraction diagonal')", True, 'import matplotlib.pyplot as plt\n'), ((62, 4, 62, 38), 'matplotlib.pyplot.ylabel', 'plt.ylabel', ({(62, 15, 62, 37): '"""log condition number"""'}, {}), "('log condition number')", True, 'import matplotlib.pyplot as plt\n'), ((70, 4, 70, 16), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((73, 4, 73, 22), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(73, 15, 73, 21): '"""N_KC"""'}, {}), "('N_KC')", True, 'import matplotlib.pyplot as plt\n'), ((79, 4, 79, 16), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((82, 4, 82, 22), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(82, 15, 82, 21): '"""N_KC"""'}, {}), "('N_KC')", True, 'import matplotlib.pyplot as plt\n'), ((84, 4, 84, 16), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((87, 4, 87, 50), 'matplotlib.pyplot.ylabel', 'plt.ylabel', ({(87, 15, 87, 49): '"""Log decrease in condition number"""'}, {}), "('Log decrease in condition number')", True, 'import matplotlib.pyplot as plt\n'), ((88, 4, 88, 22), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(88, 15, 88, 21): '"""N_KC"""'}, {}), "('N_KC')", True, 'import matplotlib.pyplot as plt\n'), ((44, 47, 44, 64), 'numpy.linalg.norm', 'np.linalg.norm', ({(44, 62, 44, 63): 'K'}, {}), '(K)', True, 'import numpy as np\n'), ((50, 19, 50, 34), 'numpy.log10', 'np.log10', ({(50, 28, 50, 33): 'conds'}, {}), '(conds)', True, 'import numpy as np\n'), ((71, 13, 71, 28), 'numpy.log10', 'np.log10', ({(71, 22, 71, 27): 'n_kcs'}, {}), '(n_kcs)', True, 'import numpy as np\n'), ((72, 15, 72, 30), 'numpy.log10', 'np.log10', ({(72, 24, 72, 29): 'n_kcs'}, {}), '(n_kcs)', True, 'import numpy as np\n'), ((80, 13, 80, 28), 'numpy.log10', 'np.log10', ({(80, 22, 80, 27): 'n_kcs'}, {}), '(n_kcs)', True, 'import numpy as np\n'), ((81, 15, 81, 30), 'numpy.log10', 'np.log10', ({(81, 24, 81, 29): 'n_kcs'}, {}), '(n_kcs)', True, 'import numpy as np\n'), ((85, 13, 85, 28), 'numpy.log10', 'np.log10', ({(85, 22, 85, 27): 'n_kcs'}, {}), '(n_kcs)', True, 'import numpy as np\n'), ((86, 15, 86, 30), 'numpy.log10', 'np.log10', ({(86, 24, 86, 29): 'n_kcs'}, {}), '(n_kcs)', True, 'import numpy as np\n'), ((27, 12, 27, 41), 'numpy.random.shuffle', 'np.random.shuffle', ({(27, 30, 27, 40): 'mask[:, (i)]'}, {}), '(mask[:, (i)])', True, 'import numpy as np\n'), ((32, 24, 32, 37), 'numpy.eye', 'np.eye', ({(32, 31, 32, 36): 'n_orn'}, {}), '(n_orn)', True, 'import numpy as np\n'), ((44, 26, 44, 43), 'numpy.linalg.pinv', 'np.linalg.pinv', ({(44, 41, 44, 42): 'K'}, {}), '(K)', True, 'import numpy as np\n'), ((67, 12, 67, 33), 'numpy.logspace', 'np.logspace', ({(67, 24, 67, 25): '1', (67, 27, 67, 28): '4', (67, 30, 67, 32): '10'}, {}), '(1, 4, 10)', True, 'import numpy as np\n'), ((76, 12, 76, 33), 'numpy.logspace', 'np.logspace', ({(76, 24, 76, 25): '1', (76, 27, 76, 28): '4', (76, 30, 76, 32): '10'}, {}), '(1, 4, 10)', True, 'import numpy as np\n')]