# python implementation of scAI optimization model
# import packages
import numpy as np
from numpy import linalg as LA
from numpy import matrix


def scAImodel_py(X1,X2,K,S=0.25,Alpha=1,Lambda=10000,Gamma=1,Maxiter = 500,Stop_rule=1,Seed=1,W1=None,W2=None,H=None,Z=None,R=None):
# parameters
    d1 = X1.shape
    d2 = X2.shape
    p = d1[0]
    n = d1[1]
    q = d2[0]
    np.random.seed(Seed)
    if W1 is None:
        W1 = np.random.rand(p,K)
        
    if W2 is None:
        W2 = np.random.rand(q,K)
    
    if H is None:
        H = np.random.rand(K,n)
    
    if Z is None:
        Z = np.random.rand(n,n)
    
    if R is None:
        R = np.random.binomial(1,S,size=(n,n))
    # main function
    XtX2 = np.dot(np.transpose(X2),X2)
    obj_old = 1
#    W1 = np.asarray(W1)
#    W2 = np.asarray(W2)
#    Z = np.asarray(Z)
#    R = np.asarray(R)
    for iter in range(1,Maxiter+1):
        #  print(iter)
        # normalize H
#        H = matrix(H)
#        lib = np.sum(H,axis = 1)
#        H = H/np.tile(lib,(1,n))
#        H = np.asarray(H)
        H = H/H.sum(axis=1,keepdims=1)
        # update W1
        HHt = np.dot(H,np.transpose(H))
        X1Ht = np.dot(X1,np.transpose(H))
        W1HHt = np.dot(W1,HHt)
        W1 = W1*X1Ht/(W1HHt+np.spacing(1))
        # update W2
        ZR = Z*R
        ZRHt = np.dot(ZR,np.transpose(H))
        X2ZRHt = np.dot(X2,ZRHt)
        W2HHt = np.dot(W2,HHt)
        W2 = W2*X2ZRHt/(W2HHt+np.spacing(1))
        # update H
        W1tX1 = np.dot(np.transpose(W1),X1)
        W2tX2 = np.dot(np.transpose(W2),X2)
        W2tX2ZR = np.dot(W2tX2,ZR)
        HZZt = np.dot(H,(Z+np.transpose(Z)))
        W1tW1 = np.dot(np.transpose(W1),W1)
        W2tW2 = np.dot(np.transpose(W2),W2)
        H = H*(Alpha*W1tX1+W2tX2ZR+Lambda*HZZt)/(np.dot(Alpha*W1tW1+W2tW2+2*Lambda*HHt+Gamma*np.ones([K,K]),H)+np.spacing(1))
        # update Z
        HtH = np.dot(np.transpose(H),H)
        X2tW2H = np.dot(np.transpose(W2tX2),H)
        RX2tW2H = R*X2tW2H
        XtX2ZR = np.dot(XtX2,ZR)
        XtX2ZRR = XtX2ZR*R
        Z = Z*(RX2tW2H+Lambda*HtH)/(XtX2ZRR+Lambda*Z+np.spacing(1))
        if Stop_rule == 2:
            obj = Alpha*pow(LA.norm(X1-np.dot(W1,H),ord = 'fro'),2)+pow(LA.norm(np.dot(X2,ZR)-np.dot(W2,H),ord = 'fro'),2)+Lambda*pow(LA.norm(Z-np.dot(np.transpose(H),H),ord = 'fro'),2)+Gamma*pow(LA.norm(np.dot(np.ones([1,K]),H),ord = 'fro'),2)
            if (obj_old-obj)/obj_old < 1e-6 and iter > 1:
                break
        iter = iter+1
            
# print ("\n")
#   print ("## Running scAI with seed %d ##" % Seed)
#   print ("\n")
    
    return W1, W2, H, Z, R