# -*- coding: utf-8 -*-
"""
Created on Thu May 28 15:54:36 2020

@author: baum_c4
"""
def sorted_stratification(X,Y,n,shuffle):
    import random
    
    import numpy as np
    import math
    
    Xl,Xn=np.shape(X)
    Yl,Yn=np.shape(Y)
    fullvector=np.concatenate([X,Y],axis=1)   
       
    sortvector=np.argsort(fullvector[:,-1])
    fullvector=fullvector[sortvector]
    
    
    
    indexes=math.floor(len(fullvector)/n)
    #splitlist=[np.array([fullvector[i+k,:] for i in range(0,indexes*10,10)]) for k in range(n) ]
    splitlist=np.array_split(fullvector, indexes)
    
    shuflist=list(range(n))
    #randi=random.sample(shuflist,len(shuflist))
    
    if shuffle==1:
        for i in range(indexes):
            randi=random.sample(shuflist,len(shuflist))
            splitlist[i]=splitlist[i][randi]
        j=[np.array([splitlist[i][k,:] for i in range(indexes)]) for k in range(n)]
    else:
        j=[np.array([splitlist[i][k,:] for i in range(indexes)]) for k in range(n)]
        #j=[np.array([fullvector[i+k,:] for i in range(0,len(fullvector)-1,n)]) for k in range(n) ]
            
    return [j[i][:,:Xn] for i in range(len(j))],[j[i][:,Xn:] for i in range(len(j))]