# To add a new cell, type '# %%'
# To add a new markdown cell, type '# %% [markdown]'
# %% [markdown]
# # 计算思维数据处理
# %% [markdown]
# 数据处理要求：
# ### 1、每个题的平均作答时长；——结果数据
# ### 2、每个题的编码种类（有多少种，分别是什么，每种多少学生）；——结果数据
# 3、每个题的每个操作步骤的平均作答时长；——过程数据
# ### 4、正确率（暂未提供标准编码，可以探索一下题目本身，协助形成标准答案编码）；——结果数据
# 5、关键节点（通过数据，探索学生在从初始状态向终止状态进行的过程中，有几个关键步骤，每个关键步骤有几种类型的关键节点编码），体现“用数据说话”去探索关键节点。——过程数据
# 6、每个题目的每种编码下，都有什么样的学生作答类型，比如都是正确的，但是可以聚成多少类，每一类有什么特征，学生是通过什么样的操作路径到达最终的。
# ### 7、每道题目的正确率；——结果数据
#
# <font color="red">注意 要求点四和要求点七相同，将一同分析
# %% [markdown]
# ## 0、数据加载和预处理
# 首先读取数据，并通过pandas进行数据帧处理
# 为了排版整洁和处理方便，以及保证工具的可拓展性，这里将处理工具封装成`data_analysis`的python类，并在jupyter中调用，该类在jupyter notebook同文件夹下的`main.py`文件中

# %%
# 首先引入需要的第三方库
import math
import pandas as pd
import json
import numpy as np
import ast
from datetime import datetime
from pandas.core import groupby
import plotly.graph_objs as go
from plotly.offline import plot
import plotly.offline as offline
from pandas.core.indexes import interval
import plotly.figure_factory as ff
pyolt=plot
import plotly.express as px
import math
import re

# %%
# 从main.py中引用对应的工具
from main import data_analysis

# %%
# 尝试将不同学校分开分析
df_all = pd.read_excel('./data/ticket_user_mianyang.xlsx')
#%%
school_list = [tup[0] for tup in list(df_all.groupby('school'))][1:]                    # 这里从index=1开始向后取元素是为了避开demo
df_list = [tup[1].reset_index(drop=True) for tup in list(df_all.groupby('school'))][1:] # 注意将index重置
df_entity_list = [data_analysis(df = df, name = school_list[i]) for i, df in enumerate(df_list)]

# %%
# 计算正确率并输出结果到output
for df_en in df_entity_list:
    df_en.calculate_acc()
    df_en.output()

# %%
# 首先将excel文件读取为pandas的dataframe类型
# 然后将该dataframe作为参数以初始化对应的数据处理工具，这里没有对excel文件进行预处理，
#   超时数据已经删除，但是并没有对不同学校进行学生分类，所以命名为'df_all'
# df_all = pd.read_excel('./data/ticket_user_mianyang.xlsx')
# dataframe并不直接对其进行处理，而是作为参数初始化一个类的实体，这样的好处是可以避免大量的代码冗余
#    在对不同学校和不同数据行进行分类后处理时，只需要额外生成新的类的实体即可
#    在这里对应所有数据行直接生成一个实体，命名为'df_all_entity'
#    注意在参数里有一个命名为'name'的参数，这里是方便在调试过程中快速判断出问题的是哪个dataframe
df_all_entity =  data_analysis(df = df_all, name = 'all')
# 该步骤运行时间较长，在32秒左右

# %%
# ## 每道题的编码种类和正确率分析
print("题目个数为：",len(df_all_entity.count_df_list))
print("参与答题总次数：",df_all_entity.row_num)

accuracy_list, addition_list = df_all_entity.calculate_acc()
prob_name_list = ['小松鼠跳柱子（1）','小松鼠跳柱子（2）','填充小石子（1）','填充小石子（2）','嵌套的矩形（1）','嵌套的矩形（2）','模式的复制（1）','模式的复制（2）','浇花（1）','浇花（2）','密码（1）','密码（2）','4进制编码（1）','4进制编码（2）','供水系统（1）','供水系统（2）', '对应的形状（1）','对应的形状（2）','滚筒（1）']
dis_list = [(prob_name_list[i], acc) for i, acc in enumerate(accuracy_list)]
dis_list

print('calculate done')

# %%
# 判断每一行中22个题目是否正确

def judge_all_data(df_all_entity, data_all, addition_list):
    columns_list = []
    for i, data in enumerate(data_all):
        data = df_all_entity.content_to_str(data)
        if i in [0,1,2,3]:
            if data in addition_list[i].index:
                columns_list.append(data)
                columns_list.append(addition_list[i].loc[data, 'success'])
            else:
                columns_list.append(data)
                columns_list.append('0')
        elif i in [5,6,7,8,18,19]:
            if data in addition_list[i-1].index:
                columns_list.append(data)
                columns_list.append(addition_list[i-1].loc[data, 'success'])
            else:
                columns_list.append(data)
                columns_list.append('0')
        elif i in [9]:
            list_temp = ast.literal_eval(data)
            if list_temp!=None and len(list_temp)==2 and type(list_temp[0])==str:
                if list_temp[0][0:2] > list_temp[0][-2:]:
                    list_temp[0] = list_temp[0][-2:] + '_' + list_temp[0][0:2]
                if list_temp[1][0:2] > list_temp[1][-2:]:
                    list_temp[1] = list_temp[1][-2:] + '_' + list_temp[1][0:2]
                if list_temp[0] > list_temp[1]:
                    list_temp =  str([list_temp[1], list_temp[0]])
                else:
                    list_temp =  str(list_temp)
            else:
                list_temp =  str(list_temp)

            columns_list.append(list_temp)
            if list_temp in addition_list[8].index:
                columns_list.append(addition_list[8].loc[list_temp, 'success'])
            else:
                columns_list.append('0')

        elif i in [10]:
            list_temp = ast.literal_eval(data)
            if type(list_temp)==list and len(list_temp)>=1 and type(list_temp[0]) == str:
                list_temp = [[int(rebuild_i) for rebuild_i in re.findall(r"\d+", str(rebuild))] for rebuild in list_temp]
                for list_mem in list_temp:
                    list_mem.sort()
                list_temp.sort()
                list_temp = str([str(list_str[0]) + '_' + str(list_str[1]) for list_str in list_temp])

            else:
                list_temp =  str(list_temp)
            columns_list.append(list_temp)
            if list_temp in addition_list[9].index:
                columns_list.append(addition_list[9].loc[list_temp, 'success'])
            else:
                columns_list.append('0')
            # columns_list.append(addition_list[i-1].loc[df_all_entity.content_to_str(data_all[i]), 'success'])
        elif i in [11]:
            list_temp = ast.literal_eval(data)
            if list_temp!=None:
                list_temp = "".join(re.findall(r"\d+", str(list_temp)))
            else:
                list_temp =  ''
            columns_list.append(list_temp)
            if list_temp in addition_list[10].index:
                columns_list.append(addition_list[10].loc[list_temp, 'success'])
            else:
                columns_list.append('0')
            # columns_list.append(addition_list[i-1].loc[df_all_entity.content_to_str(data_all[i]), 'success'])
        elif i in [12]:
            list_temp = ast.literal_eval(data)
            if list_temp!=None:
                list_temp = "".join(re.findall(r"\d+", str(list_temp)))
            else:
                list_temp =  ''

            columns_list.append(list_temp)
            if list_temp in addition_list[11].index:
                columns_list.append(addition_list[11].loc[list_temp, 'success'])
            else:
                columns_list.append('0')
            # columns_list.append(addition_list[i-1].loc[df_all_entity.content_to_str(data_all[i]), 'success'])
        elif i in [13]:
            list_temp = ast.literal_eval(data)
            if list_temp!=None:
                list_temp = "".join(re.findall(r"\d+", str(list_temp)))
            else:
                list_temp = ''
            columns_list.append(list_temp)
            if list_temp in addition_list[12].index:
                columns_list.append(addition_list[12].loc[list_temp, 'success'])
            else:
                columns_list.append('0')
            # columns_list.append(addition_list[i-1].loc[df_all_entity.content_to_str(data_all[i]), 'success'])
        elif i in [14]:
            list_temp = ast.literal_eval(data)
            if list_temp!=None:
                list_temp = "".join(re.findall(r"\d+", str(list_temp)))
            else:
                list_temp = ''
            columns_list.append(list_temp)
            if list_temp in addition_list[13].index:
                columns_list.append(addition_list[13].loc[list_temp, 'success'])
            else:
                columns_list.append('0')
            # columns_list.append(addition_list[i-1].loc[df_all_entity.content_to_str(data_all[i]), 'success'])
        elif i in [15]:
            list_temp = ast.literal_eval(data)
            if list_temp!=None:
                list_temp = "".join(re.findall(r"\d+", str(list_temp)))
            else:
                list_temp = ''
            columns_list.append(list_temp)
            if list_temp in addition_list[14].index:
                columns_list.append(addition_list[14].loc[list_temp, 'success'])
            else:
                columns_list.append('0')
            # columns_list.append(addition_list[i-1].loc[df_all_entity.content_to_str(data_all[i]), 'success'])
        elif i in [16]:
            list_temp = ast.literal_eval(data)

            if list_temp!=None:
                list_temp = "".join(re.findall(r"\d+", str(list_temp)))
            else:
                list_temp = ''
            columns_list.append(list_temp)
            if list_temp in addition_list[15].index:
                columns_list.append(addition_list[15].loc[list_temp, 'success'])
            else:
                columns_list.append('0')
            # columns_list.append(addition_list[i-1].loc[df_all_entity.content_to_str(data_all[i]), 'success'])
        elif i in [17]:
            list_temp = ast.literal_eval(data)
            if list_temp!=None:
                list_temp = "".join(re.findall(r"\d+", str(list_temp)))
            else:
                list_temp = ''
            columns_list.append(list_temp)
            if list_temp in addition_list[16].index:
                columns_list.append(addition_list[16].loc[list_temp, 'success'])
            else:
                columns_list.append('0')
            # columns_list.append(addition_list[i-1].loc[df_all_entity.content_to_str(data_all[i]), 'success'])

        elif i in [4,20,21]:
            columns_list.append(data)
            columns_list.append('unknow')

    return columns_list

df_column_prob_index = df_all_entity.df.index

df_column_prob_columns = []
for pro in range(22):
    df_column_prob_columns.append(str(pro))
    df_column_prob_columns.append('success_'+str(pro))

for df in addition_list:
    if 'list' in df.columns and 'success' in df.columns:
        for row in df.index:
            df._set_value(row,'list', str(df.loc[row, 'list']))
        df.set_index('list', inplace = True)

df_column_prob_data = []
for row in df_all_entity.df.index:
    data_all = df_all_entity.df.loc[row, 'ans']
    if len(data_all) == 22:
        df_column_prob_data.append(judge_all_data(df_all_entity, data_all, addition_list))
    else:
        df_column_prob_data.append(['None']*44)
# for row in df_all_entity.df.index:
#%%
df_column_prob = pd.DataFrame(index=df_all_entity.df.index, columns=df_column_prob_columns, data=df_column_prob_data)
df_column_prob.to_excel('./output/单个学生题目正确统计.xlsx')
df_all_entity.df.to_excel('./output/单个学生题目正确统计（学生附加信息，二者index相同）.xlsx')
pd.concat([df_all_entity.df, df_column_prob], axis=1).to_excel('./output/单个学生题目正确统计(合并).xlsx')
#%%
# 作答时间分析，在原数据上增加'interval'（做题时间）、'day'、'hour_start'、'hour_end'
data1 = list(df_all_entity.df.loc[:, 'start_time_float'])
layout={"title": "学生用时分布",
                                       "xaxis_title": "时间（24小时制）",
                                       "yaxis_title": "学生个数",
                                       # x轴坐标倾斜60度
                                       "xaxis": {"tickangle": 60}
                                      }

#数据组
hist_data=[data1]

group_labels=['做题时间分布']
import plotly.figure_factory as ff
fig=ff.create_distplot(hist_data,group_labels,bin_size=1,histnorm = 'probability')
fig['layout'].update(xaxis = dict(range = [0,24]))
plot(fig,filename='./plot/每小时回答人数统计.html')

hour_list = list(range(24))
hour_count_list = []
for hour in hour_list:
    hour_count_seq = df_all_entity.df.sort_values('start_hour').groupby('start_hour')['start_hour'].count()
    if hour not in list(hour_count_seq.index):
        hour_count_list.append(0)
    else:
        hour_count_list.append(hour_count_seq.loc[hour])

pd.DataFrame(index = hour_list, columns=['count'], data = hour_count_list).to_excel('./plot/每小时回答人数统计.xlsx')

# %%
# 正确率和编码个数绘图
## 正确的编码个数
data1 = go.Bar(x = prob_name_list, y = [len(group.groupby('success').get_group('1')) if '1' in group.groupby('success').groups.keys() else 0 for group in df_all_entity.addition_list], name = '回答中正确编码的个数')

## 错误编码的个数
data2 = go.Bar(x = prob_name_list, y = [len(group.groupby('success').get_group('0')) if '0' in group.groupby('success').groups.keys() else 0 for group in df_all_entity.addition_list], name = '回答中错误编码的个数')

## 回答正确率
data3 = go.Bar(x = prob_name_list, y = accuracy_list, name = '回答正确率')

success_count = []
for df in addition_list:
    if 'success' in df.columns:
        if '1' in df.groupby('success').groups.keys():
            success_count.append(df.groupby('success').get_group('1')['count'].sum())
        else:
            success_count.append(0)

layout={"title": "不同题目的编码数量和正确率",
       "xaxis_title": "题目编号",
       "yaxis_title": "编码个数",
       # x轴坐标倾斜60度
       "xaxis": {"tickangle": 60}
      }
from plotly.subplots import make_subplots
fig = make_subplots(rows=3, cols=1)
fig.append_trace(data1, row = 1, col = 1)
fig.append_trace(data2, row = 2, col = 1)
fig.append_trace(data3, row = 3, col = 1)

# fig = go.Figure(data=[data1, data2, data3],layout=layout)
plot(fig,filename="./plot/不同题目的编码数量和正确率.html",auto_open=False,image='png',image_height=800,image_width=1500)
pd_list = []
for i, row in enumerate(data1.y):
    pd_list.append([data1.y[i], data2.y[i], data3.y[i], success_count[i]])

pd.DataFrame(index=prob_name_list, columns=[data1.name, data2.name, data3.name, 'success_count'], data= pd_list).to_excel('./plot/不同题目的编码数量和正确率.xlsx')

#%%

data1 = go.Bar(x = school_list, y = [np.mean(df.accuracy_list) for df in df_entity_list], name = '各个学校的平均正确率')
data2 = go.Bar(x = school_list, y = [df.row_num for df in df_entity_list], name = '各个学校的参加人数')
layout={"title": "各个学校的平均正确率",
       "xaxis_title": "学校名称",
       "yaxis_title": "正确率",
       # x轴坐标倾斜60度
       "xaxis": {"tickangle": 60}
      }
from plotly.subplots import make_subplots
fig = make_subplots(rows=2, cols=1)
fig.append_trace(data1, row = 1, col = 1)
fig.append_trace(data2, row = 2, col = 1)

plot(fig,filename="./plot/各个学校的平均正确率和参加人数.html",auto_open=False,image='png',image_height=800,image_width=1500)
pd_list = []
for i, row in enumerate(data1.y):
    pd_list.append([data1.y[i], data2.y[i]])
pd.DataFrame(index=school_list, columns=[data1.name, data2.name], data= pd_list).to_excel('./plot/各个学校的平均正确率和参加人数.xlsx')
# %% [markdown]
# # 正确答案编码
# %%
# 模拟第19题的答案编码，遍历所有答案
# 所有可能的五角星（0,a）和三角形（1,b）组合
seq_list = []
for i in range(8):
    for j in range(int(math.pow(2,i+1))):
        temp=str(bin(j))[2:].zfill(i+1).replace('0','a')
        temp = temp.replace('1', 'b')
        seq_list.append(temp)
# 所有可能的长方形（1）和圆形（0）组合
trans_list = []
for i in range(3):
    for j in range(int(math.pow(2,i+1))):
        trans_list.append(str(bin(j))[2:].zfill(i+1))
# 正确序列
verify_str = '10100010010'
right_ans = []
cnt = 0
for seq in seq_list:
    for star in trans_list:
        for trian in trans_list:
            cnt +=1
            if seq.replace('a', star).replace('b', trian) == verify_str:
                right_ans.append([seq.replace('a','0').replace('b','1'), star, trian])
right_ans
# %%
# 模拟滚筒（2）的过程，遍历所有的答案
#
import math
def roll_one_time(first_dic, second_dic, verify_list):
    max_length = len(verify_list)
    line_list = [-1] * max_length
    # first roll
    for i in range(first_dic['start'], first_dic['end']+1):
        line_list[i] = first_dic['color_list'][i % len(first_dic['color_list'])]

    # second roll
    for i in range(second_dic['start'], second_dic['end']+1):
        line_list[i] = second_dic['color_list'][i % len(second_dic['color_list'])]

    is_right = True
    for i, c in enumerate(line_list):
        if c != verify_list[i]:
            is_right = False
            break
    if is_right:
        return is_right, [first_dic, second_dic]
    else:
        return is_right, []

def traverse_color(min_roll_length = 2, max_roll_length = 5, colors = [0,1,2]):
    color_lists = []
    for i in range(min_roll_length, max_roll_length+1):
        for j in range(int(math.pow(len(colors), i))):
            tmp_list = []
            for k in range(i):
                tmp_list.insert(0, colors[j % len(colors)])
                j = j // len(colors)
            color_lists.append(tmp_list)
    return color_lists

def traverse_position(start_bound = 0, end_bound = 15):
    position_list = []
    start_list = range(start_bound,end_bound+1)
    end_list = range(start_bound,end_bound+1)
    for s in start_list:
        for e in end_list:
            if e>s:
                position_list.append((s, e))
    return position_list

def traverse_roll(min_roll_length = 2, max_roll_length = 5, colors = [0,1,2], start_bound = 0, end_bound = 15):
    roll_dic_list = []
    for cl in traverse_color(min_roll_length, max_roll_length, colors):
        for p in traverse_position(start_bound, end_bound):
            if p[1]-p[0]+1>=len(cl):
                tmp_dic = {'color_list':[], 'start':0, 'end':0}
                tmp_dic['color_list']=cl
                tmp_dic['start'] = p[0]
                tmp_dic['end'] = p[1]
                roll_dic_list.append(tmp_dic)
    return roll_dic_list

# %%
# search from bottom to top

posible_roll = traverse_roll()
len(posible_roll)

# %%
posible_colors = traverse_color()
len(posible_colors)

# %%
posible_position = traverse_position()
len(posible_position)
# %%
# 从最后一次滚筒开始找起（第二次）
verify_list = [0, 1, 0, 1, 2, 2, 0, 1, 1, 2, 2, 0, 1, 1, 0, 1]
line_list = [-1]*len(verify_list)

posible_roll2_list = []
for roll2 in posible_roll:

    roll2_is_posible = True
    for index in range(roll2['start'], roll2['end']+1):
        if roll2['color_list'][((index-roll2['start']) % len(roll2['color_list']))] != verify_list[index]:
            roll2_is_posible = False
    if roll2_is_posible:
        posible_roll2_list.append(roll2)
len(posible_roll2_list)

# %%
# 根据所有可能的roll2寻找roll1
posible_roll1_roll2 = []
cnt = 0
for roll2 in posible_roll2_list:
    print(cnt)
    cnt += 1
    flag_list = [-1]*len(verify_list)
    for index2 in range(roll2['start'], roll2['end']+1):
        flag_list[index2] = -2
    for roll1 in posible_roll:
        flag_list_ = [i for i in flag_list]

        roll1_is_posible = True
        for index1 in range(roll1['start'], roll1['end']+1):
            if flag_list_[index1] == -1:
                flag_list_[index1] = -3
        if -1 in flag_list_:
            roll1_is_posible = False
            continue

        for index1 in range(roll1['start'], roll1['end']+1):
            if roll1['color_list'][((index1-roll1['start']) % len(roll1['color_list']))] != verify_list[index1] and flag_list_[index1] == -3:
                roll1_is_posible = False
                break
        if roll1_is_posible:
            posible_roll1_roll2.append((roll1, roll2))

# %%
len(posible_roll1_roll2)
# %%
posible_roll1_roll2
# %% [markdown]
# 所有可能的答案：
'''python
posible_roll1_roll2
[({'color_list': [0, 1], 'start': 0, 'end': 15},
  {'color_list': [1, 2, 2, 0, 1], 'start': 3, 'end': 12}),
 ({'color_list': [0, 1, 0, 1], 'start': 0, 'end': 15},
  {'color_list': [1, 2, 2, 0, 1], 'start': 3, 'end': 12}),
 ({'color_list': [0, 1], 'start': 0, 'end': 15},
  {'color_list': [1, 2, 2, 0, 1], 'start': 3, 'end': 13}),
 ({'color_list': [0, 1, 0, 1], 'start': 0, 'end': 15},
  {'color_list': [1, 2, 2, 0, 1], 'start': 3, 'end': 13}),
 ({'color_list': [0, 1], 'start': 0, 'end': 15},
  {'color_list': [2, 2, 0, 1, 1], 'start': 4, 'end': 12}),
 ({'color_list': [0, 1, 0, 1], 'start': 0, 'end': 15},
  {'color_list': [2, 2, 0, 1, 1], 'start': 4, 'end': 12}),
 ({'color_list': [0, 1], 'start': 0, 'end': 15},
  {'color_list': [2, 2, 0, 1, 1], 'start': 4, 'end': 13}),
 ({'color_list': [0, 1, 0, 1], 'start': 0, 'end': 15},
  {'color_list': [2, 2, 0, 1, 1], 'start': 4, 'end': 13})]
'''

#%%
# 对于滚筒（3），遍历的计算复杂度太高，这里只给出验证函数
def verify_roll_3(roll1, roll2, roll3, verify_list = [1,0,1,2,1,0,2,2,1,2,0,1,2,0,2,2,1,0,2,0]):
    flag_list = [-1]*len(verify_list)
    for index3 in range(roll3['start'], roll3['end']+1):
        if verify_list[index3] != roll3['color_list'][((index3-roll3['start']) % len(roll3['color_list']))]:
            return False
        flag_list[index3] = '3'

    for index2 in range(roll2['start'], roll2['end']+1):
        if flag_list[index2] == -1:
            if verify_list[index2] != roll2['color_list'][((index2-roll2['start']) % len(roll2['color_list']))]:
                return False
            flag_list[index2] = '2'

    for index1 in range(roll1['start'], roll1['end']+1):
        if flag_list[index1] == -1:
            if verify_list[index1] != roll1['color_list'][((index1-roll1['start']) % len(roll1['color_list']))]:
                return False
            flag_list[index1] = '1'

    if -1 in flag_list:
        return False
    else:
        return True

#%%
posible_roll = traverse_roll(min_roll_length = 2, max_roll_length = 4, colors = [0,1,2], start_bound = 0, end_bound = 19)
len(posible_roll)

#%%
# 从最后一次滚筒开始找起（第三次）
verify_list = [1,0,1,2,1,0,2,2,1,2,0,1,2,0,2,2,1,0,2,0]
line_list = [-1]*len(verify_list)

posible_roll3_list = []
for roll3 in posible_roll:

    roll3_is_posible = True
    for index in range(roll3['start'], roll3['end']+1):
        if roll3['color_list'][((index-roll3['start']) % len(roll3['color_list']))] != verify_list[index]:
            roll3_is_posible = False
    if roll3_is_posible:
        posible_roll3_list.append(roll3)
len(posible_roll3_list)

#%%
# 根据所有可能的roll3寻找roll2
posible_roll2_roll3 = []
cnt = 0
for roll3 in posible_roll3_list:
    print(cnt)
    cnt += 1
    flag_list = [-1]*len(verify_list)
    for index3 in range(roll3['start'], roll3['end']+1):
        flag_list[index3] = '3'
    for roll2 in posible_roll:
        flag_list_ = [i for i in flag_list]

        roll2_is_posible = True

        # if roll2['start']>=roll3['start'] and roll2['end']<=roll3['end']:
        #     roll2_is_posible = False
        #     continue

        for index2 in range(roll2['start'], roll2['end']+1):
            if flag_list_[index2] == -1 and roll2['color_list'][(index2-roll2['start']) % len(roll2['color_list'])] != verify_list[index2]:
                roll2_is_posible = False
                break

        if roll2_is_posible:
            posible_roll2_roll3.append((roll2, roll3))

# %%
len(posible_roll2_roll3)
# %%
verify_list = [1,0,1,2,1,0,2,2,1,2,0,1,2,0,2,2,1,0,2,0]
all_right_ans=[]
cnt = 0
for roll23 in posible_roll2_roll3:
    print(cnt,len(all_right_ans))
    cnt+=1
    for roll1 in posible_roll:
        if min([roll1['start'], roll23[0]['start'], roll23[1]['start']])!=0 or max([roll1['end'], roll23[0]['end'], roll23[1]['end']])!=len(verify_list)-1:
            continue
        if verify_roll_3(roll1, roll23[0], roll23[1], verify_list):
            all_right_ans.append((roll1, roll23[0], roll23[1]))

# %%
len(all_right_ans)
# %%
all_right_ans
# %% [markdown]
## 共有15种正确答案
'''python
[({'color_list': [1, 0], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [1, 2, 0], 'start': 8, 'end': 12}),
 ({'color_list': [1, 0, 1], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [1, 2, 0], 'start': 8, 'end': 12}),
 ({'color_list': [1, 0, 1, 0], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [1, 2, 0], 'start': 8, 'end': 12}),
 ({'color_list': [1, 0], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [1, 2, 0], 'start': 8, 'end': 13}),
 ({'color_list': [1, 0, 1], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [1, 2, 0], 'start': 8, 'end': 13}),
 ({'color_list': [1, 0, 1, 0], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [1, 2, 0], 'start': 8, 'end': 13}),
 ({'color_list': [1, 0], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [2, 0, 1], 'start': 9, 'end': 12}),
 ({'color_list': [1, 0, 1], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [2, 0, 1], 'start': 9, 'end': 12}),
 ({'color_list': [1, 0, 1, 0], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [2, 0, 1], 'start': 9, 'end': 12}),
 ({'color_list': [1, 0], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [2, 0, 1], 'start': 9, 'end': 13}),
 ({'color_list': [1, 0, 1], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [2, 0, 1], 'start': 9, 'end': 13}),
 ({'color_list': [1, 0, 1, 0], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [2, 0, 1], 'start': 9, 'end': 13}),
 ({'color_list': [1, 0], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [2, 0, 1, 2], 'start': 9, 'end': 12}),
 ({'color_list': [1, 0, 1], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [2, 0, 1, 2], 'start': 9, 'end': 12}),
 ({'color_list': [1, 0, 1, 0], 'start': 0, 'end': 19},
  {'color_list': [2, 1, 0, 2], 'start': 3, 'end': 18},
  {'color_list': [2, 0, 1, 2], 'start': 9, 'end': 12})]
'''


# #############################################################################
# misc.py
# =======
# Author : Matthieu Simeoni [matthieu.simeoni@gmail.com]
# #############################################################################

r"""
Miscellaneous functions.
"""

from typing import Tuple, Optional
import numpy as np
import re

def is_range_broadcastable(shape1: Tuple[int, int], shape2: Tuple[int, int]) -> bool:
    r"""
    Check if two shapes satisfy Numpy's broadcasting rules.

    Parameters
    ----------
    shape1: Tuple[int, int]
    shape2: Tuple[int, int]

    Returns
    -------
    bool
         ``True`` if broadcastable, ``False`` otherwise.

    Examples
    --------

    .. testsetup::

       from pycsou.util.misc import is_range_broadcastable

    .. doctest::

       >>> is_range_broadcastable((3,2), (1,2))
       True
       >>> is_range_broadcastable((3,2), (4,2))
       False
    """
    if shape1[1] != shape2[1]:
        return False
    elif shape1[0] == shape2[0]:
        return True
    elif shape1[0] == 1 or shape2[0] == 1:
        return True
    else:
        return False

def range_broadcast_shape(shape1: Tuple[int, int], shape2: Tuple[int, int]) -> Tuple[int, int]:
    r"""
    Given two shapes, determine broadcasting shape.

    Parameters
    ----------
    shape1: Tuple[int, int]
    shape2: Tuple[int, int]

    Returns
    -------
    Tuple[int, int]
        Broadcasting shape.

    Raises
    ------
    ValueError
        If the two shapes cannot be broadcasted.

    Examples
    --------

    .. testsetup::

       from pycsou.util.misc import range_broadcast_shape

    .. doctest::

       >>> range_broadcast_shape((3,2), (1,2))
       (3, 2)

    """
    if not is_range_broadcastable(shape1, shape2):
        raise ValueError('Shapes are not (range) broadcastable.')
    shape = tuple(np.fmax(shape1, shape2).tolist())
    return shape

def peaks(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    r"""
    Matlab 2D peaks function.

    Peaks is a function of two variables, obtained by translating and scaling Gaussian distributions (see `Matlab's peaks function <https://www.mathworks.com/help/matlab/ref/peaks.html>`).
    This function is useful for testing purposes.

    Parameters
    ----------
    x: np.ndarray
        X coordinates.
    y: np.ndarray
        Y coordinates.

    Returns
    -------
    np.ndarray
        Values of the 2D function ``peaks`` at the points specified by the entries of ``x`` and ``y``.

    Examples
    --------
    .. plot::

       import numpy as np
       import matplotlib.pyplot as plt
       from pycsou.util.misc import peaks

       x = np.linspace(-3,3, 1000)
       X,Y = np.meshgrid(x,x)
       Z = peaks(X,Y)
       plt.figure()
       plt.imshow(Z)

    """
    z = 3 * ((1 - x) ** 2) * np.exp(-(x ** 2) - (y + 1) ** 2) - 10 * (x / 5 - x ** 3 - y ** 5) * np.exp(
        -x ** 2 - y ** 2) - (1 / 3) * np.exp(-(x + 1) ** 2 - y ** 2)
    return z

def beamer2rst(input_file, output_file: Optional[str] = None):
    if output_file is None:
        output_file = f'{input_file.split(".")[0]}.rst'
    with open(output_file, 'w') as out_file:
        with open(input_file, 'r') as in_file:
            file_content = in_file.read()
        frames = re.findall(r'\\begin\{frame\}(.*?)\\end\{frame\}', file_content, flags=re.DOTALL)
        for frame in frames:
            frame = re.sub(r'\$\$(?P<EQ>.*?)\$\$', r'\n.. math::\n\n   \g<EQ>', frame, flags=re.DOTALL)
            frame = re.sub(r'\$(?P<EQ>.*?)\$', r':math:`\g<EQ>`', frame, flags=re.DOTALL)
            frame = re.sub(r'\\begin\{equation\}(?P<EQ>.*?)\\end\{equation\}', r'\n.. math::\n\n   \g<EQ>\n', frame,
                           flags=re.DOTALL)
            frame = re.sub(r'\\begin\{equation\*\}(?P<EQ>.*?)\\end\{equation\*\}', r'\n.. math::\n\n   \g<EQ>\n', frame,
                           flags=re.DOTALL)
            frame = re.sub(r'\\begin\{align\*\}(?P<EQ>.*?)\\end\{align\*\}', r'\n.. math::\n\n   \g<EQ>\n', frame,
                           flags=re.DOTALL)
            frame = re.sub(r'\\begin\{align\}(?P<EQ>.*?)\\end\{align\}', r'\n.. math::\n\n   \g<EQ>\n', frame,
                           flags=re.DOTALL)
            frame = re.sub(r'\\begin\{itemize\}(?P<items>.*?)\\end\{itemize\}', r'\n\g<items>\n', frame,
                           flags=re.DOTALL)
            frame = re.sub(r'\\begin\{enumerate\}(?P<items>.*?)\\end\{enumerate\}', r'\n\g<items>\n', frame,
                           flags=re.DOTALL)
            frame = re.sub(r'\\item\s*(?P<item>\S.*?)', r'* \g<item>', frame,
                           flags=re.DOTALL)
            frame = re.sub(r'\\cite\[(?P<details>.*?)\]\{(?P<citation>.*?)\}', r'\g<details> of [\g<citation>]_',
                           frame,
                           flags=re.DOTALL)
            frame = re.sub(r'\\(?P<symbol>[a-zA-Z])cal', r'\\mathcal{\g<symbol>}', frame)
            frame = re.sub(r'\\(?P<symbol>[a-zA-Z])scr', r'\\mathcal{\g<symbol>}', frame)
            frame = re.sub(r'\\(?P<symbol>[a-zA-Z])bf', r'\\mathbf{\g<symbol>}', frame)
            frame = re.sub(r'\\(?P<symbol>[a-zA-Z])bb', r'\\mathbb{\g<symbol>}', frame)
            frame = re.sub(r'\\bb(?P<symbol>[a-zA-Z])', r'\\mathbf{\g<symbol>}', frame)
            frame = re.sub(r'\\R', r'\\mathbb{R}', frame)
            frame = re.sub(r'\\N', r'\\mathbb{N}', frame)
            frame = re.sub(r'\\Q', r'\\mathbb{Q}', frame)
            frame = re.sub(r'\\bm\{(?P<symbol>.*?)\}', r'\\mathbf{\g<symbol>}', frame)
            frame = re.sub(r'\\emph\{(?P<expression>.*?)\}', r'*\g<expression>*', frame)
            frame = re.sub(r'\\textbf\{(?P<expression>.*?)\}', r'**\g<expression>**', frame)
            frame = re.sub(r'(\\green|\\blue|\\red|\\orange|\\purple)\{(?P<expression>.*?)\}', r'\g<expression>',
                           frame)
            out_file.write(frame)

if __name__ == '__main__':
    beamer2rst('pycsou/util/part1.tex')


import pytest
import matplotlib.pyplot as plt
import numpy as np

# The following tests work for all three themes as these tests are scanning through the elements within the themes
# but the framework of three themes remain the same.

# Creating random data for test plot
N = 500
x = np.random.rand(N)
y = np.random.rand(N)
z = np.random.rand(N)
colors = (0,0,0)
area = np.pi*3

# Plotting a scatterplot with parameters to test
plt.scatter(x, y, c="red", alpha=0.5, label = "y")
plt.scatter(x, z, c="blue", alpha=0.5, label = "z")
plt.legend(loc="upper right",title = "hi")
plt.rcParams['axes.facecolor']= "white"
plt.rcParams['axes.titlesize'] = 20
plt.rcParams['axes.edgecolor'] = "black"
plt.rcParams['legend.fontsize'] = 12
plt.rcParams["legend.title_fontsize"] = 12
plt.rcParams["legend.facecolor"] = "white"
plt.rcParams["figure.facecolor"] = "green"
plt.rcParams["axes.labelsize"] = 20
plt.rcParams["scatter.marker"] = "*"
plt.title('Testing Graph')
plt.xlabel("X")
plt.ylabel("Y")
plt.show()
plt.close()

def test_facecolor():
    ''' A function that checks the face colour of the plot background  '''

    assert plt.rcParams['axes.facecolor'] == "white"

def test_title_size():
    '''A function that checks that the title font size of a graph '''

    assert plt.rcParams['axes.titlesize'] == 20

def test_legend_fontsize():
    ''''A function that checks that the font size of the legend title of a graph  '''

    assert plt.rcParams['legend.fontsize'] == 12

def test_legend_facecolor():
    ''' 'A function that checks the legend face colour of a graph '''

    assert plt.rcParams["legend.facecolor"] == "white"

def test_figure_facecolor():
    ''' A function that checks the plot's surrounding colour of a graph '''

    assert plt.rcParams["figure.facecolor"] == "green"

def test_axes_labelsize():
    ''' A function that checks the label fond size of the axis titles'''

    assert plt.rcParams["axes.labelsize"] == 20

def test_scatter_marker():
    '''A function that checks the shape of the scatter points '''

    assert plt.rcParams["scatter.marker"] == "*"

# Fails expected

def test_axes_edgecolor():
    ''' A function that checks the edge colours of the axis  '''

    assert plt.rcParams['axes.edgecolor'] == "blue"

def test_legend_fontsize_neg():
    ''' A function that tests the label font size of the axis titles, that should fail'''

    assert plt.rcParams['axes.labelsize'] == 1

def test_legend_title_fontsize():
    ''' A test that checks the title font size of legend '''

    assert plt.rcParams["legend.title_fontsize"] == 45

def test_title_size_neg():
    '''A test that checks that the title font size of a graph, that should fail  '''

    assert plt.rcParams['axes.titlesize'] == 18


#!/usr/bin/env python
# coding: utf-8

# # Project 3: Smart Beta Portfolio and Portfolio Optimization
#
# ## Overview
#
#
# Smart beta has a broad meaning, but we can say in practice that when we use the universe of stocks from an index, and then apply some weighting scheme other than market cap weighting, it can be considered a type of smart beta fund.  A Smart Beta portfolio generally gives investors exposure or "beta" to one or more types of market characteristics (or factors) that are believed to predict prices while giving investors a diversified broad exposure to a particular market. Smart Beta portfolios generally target momentum, earnings quality, low volatility, and dividends or some combination. Smart Beta Portfolios are generally rebalanced infrequently and follow relatively simple rules or algorithms that are passively managed.  Model changes to these types of funds are also rare requiring prospectus filings with US Security and Exchange Commission in the case of US focused mutual funds or ETFs.. Smart Beta portfolios are generally long-only, they do not short stocks.
#
# In contrast, a purely alpha-focused quantitative fund may use multiple models or algorithms to create a portfolio. The portfolio manager retains discretion in upgrading or changing the types of models and how often to rebalance the portfolio in attempt to maximize performance in comparison to a stock benchmark.  Managers may have discretion to short stocks in portfolios.
#
# Imagine you're a portfolio manager, and wish to try out some different portfolio weighting methods.
#
# One way to design portfolio is to look at certain accounting measures (fundamentals) that, based on past trends, indicate stocks that produce better results.
#
#
# For instance, you may start with a hypothesis that dividend-issuing stocks tend to perform better than stocks that do not. This may not always be true of all companies; for instance, Apple does not issue dividends, but has had good historical performance.  The hypothesis about dividend-paying stocks may go something like this:
#
# Companies that regularly issue dividends may also be more prudent in allocating their available cash, and may indicate that they are more conscious of prioritizing shareholder interests.  For example, a CEO may decide to reinvest cash into pet projects that produce low returns.  Or, the CEO may do some analysis, identify that reinvesting within the company produces lower returns compared to a diversified portfolio, and so decide that shareholders would be better served if they were given the cash (in the form of dividends).  So according to this hypothesis, dividends may be both a proxy for how the company is doing (in terms of earnings and cash flow), but also a signal that the company acts in the best interest of its shareholders.  Of course, it's important to test whether this works in practice.
#
#
# You may also have another hypothesis, with which you wish to design a portfolio that can then be made into an ETF.  You may find that investors may wish to invest in passive beta funds, but wish to have less risk exposure (less volatility) in their investments.  The goal of having a low volatility fund that still produces returns similar to an index may be appealing to investors who have a shorter investment time horizon, and so are more risk averse.
#
# So the objective of your proposed portfolio is to design a portfolio that closely tracks an index, while also minimizing the portfolio variance.  Also, if this portfolio can match the returns of the index with less volatility, then it has a higher risk-adjusted return (same return, lower volatility).
#
# Smart Beta ETFs can be designed with both of these two general methods (among others): alternative weighting and minimum volatility ETF.
#
#
# ## Instructions
# Each problem consists of a function to implement and instructions on how to implement the function.  The parts of the function that need to be implemented are marked with a `# TODO` comment. After implementing the function, run the cell to test it against the unit tests we've provided. For each problem, we provide one or more unit tests from our `project_tests` package. These unit tests won't tell you if your answer is correct, but will warn you of any major errors. Your code will be checked for the correct solution when you submit it to Udacity.
#
# ## Packages
# When you implement the functions, you'll only need to you use the packages you've used in the classroom, like [Pandas](https://pandas.pydata.org/) and [Numpy](http://www.numpy.org/). These packages will be imported for you. We recommend you don't add any import statements, otherwise the grader might not be able to run your code.
#
# The other packages that we're importing are `helper`, `project_helper`, and `project_tests`. These are custom packages built to help you solve the problems.  The `helper` and `project_helper` module contains utility functions and graph functions. The `project_tests` contains the unit tests for all the problems.
# ### Install Packages

# In[1]:

import sys
get_ipython().system('{sys.executable} -m pip install -r requirements.txt')

# In[2]:

import pandas as pd
import numpy as np
import helper
import project_helper
import project_tests

#
# ### Load Packages

# ## Market Data
# ### Load Data
# For this universe of stocks, we'll be selecting large dollar volume stocks. We're using this universe, since it is highly liquid.

# In[3]:

df = pd.read_csv('../../data/project_3/eod-quotemedia.csv')

percent_top_dollar = 0.2
high_volume_symbols = project_helper.large_dollar_volume_stocks(df, 'adj_close', 'adj_volume', percent_top_dollar)
df = df[df['ticker'].isin(high_volume_symbols)]

close = df.reset_index().pivot(index='date', columns='ticker', values='adj_close')
volume = df.reset_index().pivot(index='date', columns='ticker', values='adj_volume')
dividends = df.reset_index().pivot(index='date', columns='ticker', values='dividends')

# ### View Data
# To see what one of these 2-d matrices looks like, let's take a look at the closing prices matrix.

# In[4]:

project_helper.print_dataframe(close)

# # Part 1: Smart Beta Portfolio
# In Part 1 of this project, you'll build a portfolio using dividend yield to choose the portfolio weights. A portfolio such as this could be incorporated into a smart beta ETF.  You'll compare this portfolio to a market cap weighted index to see how well it performs.
#
# Note that in practice, you'll probably get the index weights from a data vendor (such as companies that create indices, like MSCI, FTSE, Standard and Poor's), but for this exercise we will simulate a market cap weighted index.
#
# ## Index Weights
# The index we'll be using is based on large dollar volume stocks. Implement `generate_dollar_volume_weights` to generate the weights for this index. For each date, generate the weights based on dollar volume traded for that date. For example, assume the following is close prices and volume data:
# ```
#                  Prices
#                A         B         ...
# 2013-07-08     2         2         ...
# 2013-07-09     5         6         ...
# 2013-07-10     1         2         ...
# 2013-07-11     6         5         ...
# ...            ...       ...       ...
#
#                  Volume
#                A         B         ...
# 2013-07-08     100       340       ...
# 2013-07-09     240       220       ...
# 2013-07-10     120       500       ...
# 2013-07-11     10        100       ...
# ...            ...       ...       ...
# ```
# The weights created from the function `generate_dollar_volume_weights` should be the following:
# ```
#                A         B         ...
# 2013-07-08     0.126..   0.194..   ...
# 2013-07-09     0.759..   0.377..   ...
# 2013-07-10     0.075..   0.285..   ...
# 2013-07-11     0.037..   0.142..   ...
# ...            ...       ...       ...
# ```

# In[5]:

def generate_dollar_volume_weights(close, volume):
    """
    Generate dollar volume weights.

    Parameters
    ----------
    close : DataFrame
        Close price for each ticker and date
    volume : str
        Volume for each ticker and date

    Returns
    -------
    dollar_volume_weights : DataFrame
        The dollar volume weights for each ticker and date
    """

    product = close*volume

    weights=product.apply(lambda r : r/sum(r),axis=1)

    assert close.index.equals(volume.index)
    assert close.columns.equals(volume.columns)

    #TODO: Implement function

    return weights

project_tests.test_generate_dollar_volume_weights(generate_dollar_volume_weights)

# ### View Data
# Let's generate the index weights using `generate_dollar_volume_weights` and view them using a heatmap.

# In[6]:

index_weights = generate_dollar_volume_weights(close, volume)
project_helper.plot_weights(index_weights, 'Index Weights')

# ## Portfolio Weights
# Now that we have the index weights, let's choose the portfolio weights based on dividend. You would normally calculate the weights based on trailing dividend yield, but we'll simplify this by just calculating the total dividend yield over time.
#
# Implement `calculate_dividend_weights` to return the weights for each stock based on its total dividend yield over time. This is similar to generating the weight for the index, but it's using dividend data instead.
# For example, assume the following is `dividends` data:
# ```
#                  Prices
#                A         B
# 2013-07-08     0         0
# 2013-07-09     0         1
# 2013-07-10     0.5       0
# 2013-07-11     0         0
# 2013-07-12     2         0
# ...            ...       ...
# ```
# The weights created from the function `calculate_dividend_weights` should be the following:
# ```
#                A         B
# 2013-07-08     NaN       NaN
# 2013-07-09     0         1
# 2013-07-10     0.333..   0.666..
# 2013-07-11     0.333..   0.666..
# 2013-07-12     0.714..   0.285..
# ...            ...       ...
# ```

# In[7]:

def calculate_dividend_weights(dividends):

    """
    Calculate dividend weights.

    Parameters
    ----------
    dividends : DataFrame
        Dividend for each stock and date

    Returns
    -------
    dividend_weights : DataFrame
        Weights for each stock and date
    """
    #TODO: Implement function
    divs=dividends.copy()

    weights_r=divs.cumsum()

    weights=weights_r.apply((lambda x: x/sum(x)),axis=1)
    return weights

project_tests.test_calculate_dividend_weights(calculate_dividend_weights)

# ### View Data
# Just like the index weights, let's generate the ETF weights and view them using a heatmap.

# In[8]:

etf_weights = calculate_dividend_weights(dividends)
project_helper.plot_weights(etf_weights, 'ETF Weights')

# ## Returns
# Implement `generate_returns` to generate returns data for all the stocks and dates from price data. You might notice we're implementing returns and not log returns. Since we're not dealing with volatility, we don't have to use log returns.

# In[9]:

def generate_returns(prices):
    """
    Generate returns for ticker and date.

    Parameters
    ----------
    prices : DataFrame
        Price for each ticker and date

    Returns
    -------
    returns : Dataframe
        The returns for each ticker and date
    """
    r=prices.copy()

    d=r.shift(1)

    returns=(r-d)/d

    #TODO: Implement function

    return returns

project_tests.test_generate_returns(generate_returns)

# ### View Data
# Let's generate the closing returns using `generate_returns` and view them using a heatmap.

# In[10]:

returns = generate_returns(close)
project_helper.plot_returns(returns, 'Close Returns')

# ## Weighted Returns
# With the returns of each stock computed, we can use it to compute the returns for an index or ETF. Implement `generate_weighted_returns` to create weighted returns using the returns and weights.

# In[11]:

def generate_weighted_returns(returns, weights):
    """
    Generate weighted returns.

    Parameters
    ----------
    returns : DataFrame
        Returns for each ticker and date
    weights : DataFrame
        Weights for each ticker and date

    Returns
    -------
    weighted_returns : DataFrame
        Weighted returns for each ticker and date
    """
    weighted_returns=returns*weights

    assert returns.index.equals(weights.index)
    assert returns.columns.equals(weights.columns)

    #TODO: Implement function

    return weighted_returns

project_tests.test_generate_weighted_returns(generate_weighted_returns)

# ### View Data
# Let's generate the ETF and index returns using `generate_weighted_returns` and view them using a heatmap.

# In[12]:

index_weighted_returns = generate_weighted_returns(returns, index_weights)
etf_weighted_returns = generate_weighted_returns(returns, etf_weights)
project_helper.plot_returns(index_weighted_returns, 'Index Returns')
project_helper.plot_returns(etf_weighted_returns, 'ETF Returns')

# ## Cumulative Returns
# To compare performance between the ETF and Index, we're going to calculate the tracking error. Before we do that, we first need to calculate the index and ETF comulative returns. Implement `calculate_cumulative_returns` to calculate the cumulative returns over time given the returns.

# In[13]:

def calculate_cumulative_returns(returns):
    """
    Calculate cumulative returns.

    Parameters
    ----------
    returns : DataFrame
        Returns for each ticker and date

    Returns
    -------
    cumulative_returns : Pandas Series
        Cumulative returns for each date
    """
    r=returns.sum(axis=1)
    ret=r.apply(lambda x: x+1)
    cum_ret=ret.cumprod()
    #TODO: Implement function

    return cum_ret

project_tests.test_calculate_cumulative_returns(calculate_cumulative_returns)

# ### View Data
# Let's generate the ETF and index cumulative returns using `calculate_cumulative_returns` and compare the two.

# In[14]:

index_weighted_cumulative_returns = calculate_cumulative_returns(index_weighted_returns)
etf_weighted_cumulative_returns = calculate_cumulative_returns(etf_weighted_returns)
project_helper.plot_benchmark_returns(index_weighted_cumulative_returns, etf_weighted_cumulative_returns, 'Smart Beta ETF vs Index')

# ## Tracking Error
# In order to check the performance of the smart beta portfolio, we can calculate the annualized tracking error against the index. Implement `tracking_error` to return the tracking error between the ETF and benchmark.
#
# For reference, we'll be using the following annualized tracking error function:
# $$ TE = \sqrt{252} * SampleStdev(r_p - r_b) $$
#
# Where $ r_p $ is the portfolio/ETF returns and $ r_b $ is the benchmark returns.
#
# _Note: When calculating the sample standard deviation, the delta degrees of freedom is 1, which is the also the default value._

# In[15]:

def tracking_error(benchmark_returns_by_date, etf_returns_by_date):
    """
    Calculate the tracking error.

    Parameters
    ----------
    benchmark_returns_by_date : Pandas Series
        The benchmark returns for each date
    etf_returns_by_date : Pandas Series
        The ETF returns for each date

    Returns
    -------
    tracking_error : float
        The tracking error
    """
    series_ret=(etf_returns_by_date)-(benchmark_returns_by_date)
    stdev=series_ret.std()
    Tracking_error=np.sqrt(252)*stdev
    assert benchmark_returns_by_date.index.equals(etf_returns_by_date.index)

    #TODO: Implement function

    return Tracking_error

project_tests.test_tracking_error(tracking_error)

# ### View Data
# Let's generate the tracking error using `tracking_error`.

# In[16]:

smart_beta_tracking_error = tracking_error(np.sum(index_weighted_returns, 1), np.sum(etf_weighted_returns, 1))
print('Smart Beta Tracking Error: {}'.format(smart_beta_tracking_error))

# # Part 2: Portfolio Optimization
#
# Now, let's create a second portfolio.  We'll still reuse the market cap weighted index, but this will be independent of the dividend-weighted portfolio that we created in part 1.
#
# We want to both minimize the portfolio variance and also want to closely track a market cap weighted index.  In other words, we're trying to minimize the distance between the weights of our portfolio and the weights of the index.
#
# $Minimize \left [ \sigma^2_p + \lambda \sqrt{\sum_{1}^{m}(weight_i - indexWeight_i)^2} \right  ]$ where $m$ is the number of stocks in the portfolio, and $\lambda$ is a scaling factor that you can choose.
#
# Why are we doing this? One way that investors evaluate a fund is by how well it tracks its index. The fund is still expected to deviate from the index within a certain range in order to improve fund performance.  A way for a fund to track the performance of its benchmark is by keeping its asset weights similar to the weights of the index.  We’d expect that if the fund has the same stocks as the benchmark, and also the same weights for each stock as the benchmark, the fund would yield about the same returns as the benchmark. By minimizing a linear combination of both the portfolio risk and distance between portfolio and benchmark weights, we attempt to balance the desire to minimize portfolio variance with the goal of tracking the index.
#
#
# ## Covariance
# Implement `get_covariance_returns` to calculate the covariance of the `returns`. We'll use this to calculate the portfolio variance.
#
# If we have $m$ stock series, the covariance matrix is an $m \times m$ matrix containing the covariance between each pair of stocks.  We can use [`Numpy.cov`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.cov.html) to get the covariance.  We give it a 2D array in which each row is a stock series, and each column is an observation at the same period of time. For any `NaN` values, you can replace them with zeros using the [`DataFrame.fillna`](https://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.fillna.html) function.
#
# The covariance matrix $\mathbf{P} =
# \begin{bmatrix}
# \sigma^2_{1,1} & ... & \sigma^2_{1,m} \\
# ... & ... & ...\\
# \sigma_{m,1} & ... & \sigma^2_{m,m}  \\
# \end{bmatrix}$

# In[17]:

def get_covariance_returns(returns):
    """
    Calculate covariance matrices.

    Parameters
    ----------
    returns : DataFrame
        Returns for each ticker and date

    Returns
    -------
    returns_covariance  : 2 dimensional Ndarray
        The covariance of the returns
    """
    #TODO: Implement function

    cov=returns.fillna(0)

    return np.cov(cov,rowvar=False)

project_tests.test_get_covariance_returns(get_covariance_returns)

# ### View Data
# Let's look at the covariance generated from `get_covariance_returns`.

# In[18]:

covariance_returns = get_covariance_returns(returns)
covariance_returns = pd.DataFrame(covariance_returns, returns.columns, returns.columns)

covariance_returns_correlation = np.linalg.inv(np.diag(np.sqrt(np.diag(covariance_returns))))
covariance_returns_correlation = pd.DataFrame(
    covariance_returns_correlation.dot(covariance_returns).dot(covariance_returns_correlation),
    covariance_returns.index,
    covariance_returns.columns)

project_helper.plot_covariance_returns_correlation(
    covariance_returns_correlation,
    'Covariance Returns Correlation Matrix')

# ### portfolio variance
# We can write the portfolio variance $\sigma^2_p = \mathbf{x^T} \mathbf{P} \mathbf{x}$
#
# Recall that the $\mathbf{x^T} \mathbf{P} \mathbf{x}$ is called the quadratic form.
# We can use the cvxpy function `quad_form(x,P)` to get the quadratic form.
#
# ### Distance from index weights
# We want portfolio weights that track the index closely.  So we want to minimize the distance between them.
# Recall from the Pythagorean theorem that you can get the distance between two points in an x,y plane by adding the square of the x and y distances and taking the square root.  Extending this to any number of dimensions is called the L2 norm.  So: $\sqrt{\sum_{1}^{n}(weight_i - indexWeight_i)^2}$  Can also be written as $\left \| \mathbf{x} - \mathbf{index} \right \|_2$.  There's a cvxpy function called [norm()](https://www.cvxpy.org/api_reference/cvxpy.atoms.other_atoms.html#norm)
# `norm(x, p=2, axis=None)`.  The default is already set to find an L2 norm, so you would pass in one argument, which is the difference between your portfolio weights and the index weights.
#
# ### objective function
# We want to minimize both the portfolio variance and the distance of the portfolio weights from the index weights.
# We also want to choose a `scale` constant, which is $\lambda$ in the expression.
#
# $\mathbf{x^T} \mathbf{P} \mathbf{x} + \lambda \left \| \mathbf{x} - \mathbf{index} \right \|_2$
#
#
# This lets us choose how much priority we give to minimizing the difference from the index, relative to minimizing the variance of the portfolio.  If you choose a higher value for `scale` ($\lambda$).
#
# We can find the objective function using cvxpy `objective = cvx.Minimize()`.  Can you guess what to pass into this function?
#
#

# ### constraints
# We can also define our constraints in a list.  For example, you'd want the weights to sum to one. So $\sum_{1}^{n}x = 1$.  You may also need to go long only, which means no shorting, so no negative weights.  So $x_i >0 $ for all $i$. you could save a variable as `[x >= 0, sum(x) == 1]`, where x was created using `cvx.Variable()`.
#
# ### optimization
# So now that we have our objective function and constraints, we can solve for the values of $\mathbf{x}$.
# cvxpy has the constructor `Problem(objective, constraints)`, which returns a `Problem` object.
#
# The `Problem` object has a function solve(), which returns the minimum of the solution.  In this case, this is the minimum variance of the portfolio.
#
# It also updates the vector $\mathbf{x}$.
#
# We can check out the values of $x_A$ and $x_B$ that gave the minimum portfolio variance by using `x.value`

# In[19]:

import cvxpy as cvx

def get_optimal_weights(covariance_returns, index_weights, scale=2.0):
    """
    Find the optimal weights.

    Parameters
    ----------
    covariance_returns : 2 dimensional Ndarray
        The covariance of the returns
    index_weights : Pandas Series
        Index weights for all tickers at a period in time
    scale : int
        The penalty factor for weights the deviate from the index
    Returns
    -------
    x : 1 dimensional Ndarray
        The solution for x
    """

    m=index_weights.shape[0]

    x=cvx.Variable(m)

    p = covariance_returns

    portfolio_variance = cvx.quad_form(x,p)

    diff=(x-index_weights)

    dist=cvx.norm(diff,p=2,axis=None)

    objective = cvx.Minimize(portfolio_variance + scale*(dist))

    constraints = [x >= 0 , sum(x) == 1]

    problem= cvx.Problem(objective,constraints)

    solved= problem.solve()

    x_vals= x.value

    assert len(covariance_returns.shape) == 2
    assert len(index_weights.shape) == 1
    assert covariance_returns.shape[0] == covariance_returns.shape[1]  == index_weights.shape[0]

    #TODO: Implement function

    return x_vals

project_tests.test_get_optimal_weights(get_optimal_weights)

# ## Optimized Portfolio
# Using the `get_optimal_weights` function, let's generate the optimal ETF weights without rebalanceing. We can do this by feeding in the covariance of the entire history of data. We also need to feed in a set of index weights. We'll go with the average weights of the index over time.

# In[20]:

raw_optimal_single_rebalance_etf_weights = get_optimal_weights(covariance_returns.values, index_weights.iloc[-1])
optimal_single_rebalance_etf_weights = pd.DataFrame(
    np.tile(raw_optimal_single_rebalance_etf_weights, (len(returns.index), 1)),
    returns.index,
    returns.columns)

# With our ETF weights built, let's compare it to the index. Run the next cell to calculate the ETF returns and compare it to the index returns.

# In[21]:

optim_etf_returns = generate_weighted_returns(returns, optimal_single_rebalance_etf_weights)
optim_etf_cumulative_returns = calculate_cumulative_returns(optim_etf_returns)
project_helper.plot_benchmark_returns(index_weighted_cumulative_returns, optim_etf_cumulative_returns, 'Optimized ETF vs Index')

optim_etf_tracking_error = tracking_error(np.sum(index_weighted_returns, 1), np.sum(optim_etf_returns, 1))
print('Optimized ETF Tracking Error: {}'.format(optim_etf_tracking_error))

# ## Rebalance Portfolio Over Time
# The single optimized ETF portfolio used the same weights for the entire history. This might not be the optimal weights for the entire period. Let's rebalance the portfolio over the same period instead of using the same weights. Implement `rebalance_portfolio` to rebalance a portfolio.
#
# Reblance the portfolio every n number of days, which is given as `shift_size`. When rebalancing, you should look back a certain number of days of data in the past, denoted as `chunk_size`. Using this data, compute the optoimal weights using `get_optimal_weights` and `get_covariance_returns`.

# In[22]:

def rebalance_portfolio(returns, index_weights, shift_size, chunk_size):
    """
    Get weights for each rebalancing of the portfolio.

    Parameters
    ----------
    returns : DataFrame
        Returns for each ticker and date
    index_weights : DataFrame
        Index weight for each ticker and date
    shift_size : int
        The number of days between each rebalance
    chunk_size : int
        The number of days to look in the past for rebalancing

    Returns
    -------
    all_rebalance_weights  : list of Ndarrays
        The ETF weights for each point they are rebalanced
    """
    returns = returns.fillna(0).copy()

    generate=np.arange(1,len(returns))

    lists_i=[]

    count=0

    lists_r=[]

    count=-1
    for i in generate:
        count+=1
        if len(lists_r)!=chunk_size:
            lists_r.append(count)

        if len(lists_r)==chunk_size:
            n = lists_r[shift_size-chunk_size:]
            count=n[-1]
            lists_i.append(lists_r)
            lists_r = []
            for i in n:
                lists_r.append(i)
    lists_d=[]

    for d in lists_i :
        lists_d.append(d[-1])

    list_vals=[]
    for i in lists_i:
        list_input=[]
        for n in i:
            list_input.append(returns.iloc[n].values)

        list_vals.append(list_input)
        list_input=[]

    v=np.array(list_vals)

    rest = []
    for i in v:
        k= pd.DataFrame(i)
        rest.append(get_covariance_returns(k))

    weights = []

    for (i,n) in zip(lists_d,rest):
        index_w = index_weights.iloc[i]
        weights.append(get_optimal_weights(n,index_w,scale=2.0))

    assert returns.index.equals(index_weights.index)
    assert returns.columns.equals(index_weights.columns)
    assert shift_size > 0
    assert chunk_size >= 0

    #TODO: Implement function

    return weights

project_tests.test_rebalance_portfolio(rebalance_portfolio)

# Run the following cell to rebalance the portfolio using `rebalance_portfolio`.

# In[23]:

chunk_size = 250
shift_size = 5
all_rebalance_weights = rebalance_portfolio(returns, index_weights, shift_size, chunk_size)

# ## Portfolio Turnover
# With the portfolio rebalanced, we need to use a metric to measure the cost of rebalancing the portfolio. Implement `get_portfolio_turnover` to calculate the annual portfolio turnover. We'll be using the formulas used in the classroom:
#
# $ AnnualizedTurnover =\frac{SumTotalTurnover}{NumberOfRebalanceEvents} * NumberofRebalanceEventsPerYear $
#
# $ SumTotalTurnover =\sum_{t,n}{\left | x_{t,n} - x_{t+1,n} \right |} $ Where $ x_{t,n} $ are the weights at time $ t $ for equity $ n $.
#
# $ SumTotalTurnover $ is just a different way of writing $ \sum \left | x_{t_1,n} - x_{t_2,n} \right | $

# In[40]:

def get_portfolio_turnover(all_rebalance_weights, shift_size, rebalance_count, n_trading_days_in_year=252):
    """
    Calculage portfolio turnover.

    Parameters
    ----------
    all_rebalance_weights : list of Ndarrays
        The ETF weights for each point they are rebalanced
    shift_size : int
        The number of days between each rebalance
    rebalance_count : int
        Number of times the portfolio was rebalanced
    n_trading_days_in_year: int
        Number of trading days in a year

    Returns
    -------
    portfolio_turnover  : float
        The portfolio turnover
    """
    rebalance_e = n_trading_days_in_year/shift_size

    df = pd.DataFrame(all_rebalance_weights)
    df_s=df.diff(periods=-1)
    df_t=df_s.fillna(0)

    turnover_per_n = df_t.apply( lambda x: abs(x)).sum()

    sum_total_turnover = turnover_per_n.sum()

    annualized_turnover = (sum_total_turnover*rebalance_e)/rebalance_count

    assert shift_size > 0
    assert rebalance_count > 0

    #TODO: Implement function

    return annualized_turnover

project_tests.test_get_portfolio_turnover(get_portfolio_turnover)

# Run the following cell to get the portfolio turnover from  `get_portfolio turnover`.

# In[41]:

print(get_portfolio_turnover(all_rebalance_weights, shift_size, len(all_rebalance_weights) - 1))

# That's it! You've built a smart beta portfolio in part 1 and did portfolio optimization in part 2. You can now submit your project.

# ## Submission
# Now that you're done with the project, it's time to submit it. Click the submit button in the bottom right. One of our reviewers will give you feedback on your project with a pass or not passed grade. You can continue to the next section while you wait for feedback.


#twoDimensionalListAndNumpyArray.py
import numpy as np

twoDimList = [[1,2,3,4],[2,3,4,5]]
print(twoDimList)
twoDimArray = np.array(twoDimList)
print(twoDimArray)

for i in range(len(twoDimList)):
    print(twoDimList[i])
    print(twoDimArray[i])
    for j in range(len(twoDimList[i])):
        print(twoDimList[i][j])
        print(twoDimArray[i][j])


# -------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
# -------------------------------------------------------------

import io
import json
import os
import shutil
import sys
import unittest

import numpy as np
from systemds.context import SystemDSContext
from systemds.matrix import Federated

os.environ['SYSDS_QUIET'] = "1"

dim = 5
np.random.seed(132)
m1 = np.array(np.random.randint(100, size=dim * dim) + 1.01, dtype=np.double)
m1.shape = (dim, dim)
m2 = np.array(np.random.randint(5, size=dim * dim) + 1, dtype=np.double)
m2.shape = (dim, dim)

tempdir = "/tmp/test_federated_basic/"
mtd = {"format": "csv", "header": "false", "rows": dim, "cols": dim}

# Create the testing directory if it does not exist.
if not os.path.exists(tempdir):
    os.makedirs(tempdir)

# Save data files for the Federated workers.
np.savetxt(tempdir + "m1.csv", m1, delimiter=",")
with io.open(tempdir + "m1.csv.mtd", "w", encoding="utf-8") as f:
    f.write(json.dumps(mtd, ensure_ascii=False))

np.savetxt(tempdir + "m2.csv", m2, delimiter=",")
with io.open(tempdir + "m2.csv.mtd", "w", encoding="utf-8") as f:
    f.write(json.dumps(mtd, ensure_ascii=False))

# Federated workers + file locations
fed1 = "localhost:8001/" + tempdir + "m1.csv"
fed2 = "localhost:8002/" + tempdir + "m2.csv"

class TestFederatedAggFn(unittest.TestCase):

    sds: SystemDSContext = None

    @classmethod
    def setUpClass(cls):
        cls.sds = SystemDSContext()

    @classmethod
    def tearDownClass(cls):
        cls.sds.close()

    def test_1(self):
        f_m1 = Federated(self.sds, [fed1], [([0, 0], [dim, dim])]).compute()
        res = np.allclose(f_m1, m1)
        self.assertTrue(res)

    def test_2(self):
        f_m2 = Federated(self.sds, [fed2], [([0, 0], [dim, dim])]).compute()
        res = np.allclose(f_m2, m2)
        self.assertTrue(res)

    def test_3(self):
        #   [[m1,m1,m1,m1,m1,m2,m2,m2,m2,m2]
        #    [m1,m1,m1,m1,m1,m2,m2,m2,m2,m2]
        #    [m1,m1,m1,m1,m1,m2,m2,m2,m2,m2]
        #    [m1,m1,m1,m1,m1,m2,m2,m2,m2,m2]
        #    [m1,m1,m1,m1,m1,m2,m2,m2,m2,m2]]
        f_m1_m2 = Federated(self.sds,
            [fed1, fed2],
            [([0, 0], [dim, dim]), ([0, dim], [dim, dim * 2])]
        ).compute()
        m1_m2 = np.concatenate((m1, m2), axis=1)
        res = np.allclose(f_m1_m2, m1_m2)
        self.assertTrue(res)

    def test_4(self):
        #   [[m1,m1,m1,m1,m1]
        #    [m1,m1,m1,m1,m1]
        #    [m1,m1,m1,m1,m1]
        #    [m1,m1,m1,m1,m1]
        #    [m1,m1,m1,m1,m1]
        #    [m2,m2,m2,m2,m2]
        #    [m2,m2,m2,m2,m2]
        #    [m2,m2,m2,m2,m2]
        #    [m2,m2,m2,m2,m2]
        #    [m2,m2,m2,m2,m2]]
        f_m1_m2 = Federated(self.sds,
            [fed1, fed2],
            [([0, 0], [dim, dim]), ([dim, 0], [dim * 2, dim])]
        ).compute()
        m1_m2 = np.concatenate((m1, m2))
        res = np.allclose(f_m1_m2, m1_m2)
        self.assertTrue(res)

    # -----------------------------------
    # The rest of the tests are
    # Extended functionality not working Yet
    # -----------------------------------

    def test_5(self):
        #   [[m1,m1,m1,m1,m1, 0, 0, 0, 0, 0]
        #    [m1,m1,m1,m1,m1, 0, 0, 0, 0, 0]
        #    [m1,m1,m1,m1,m1,m2,m2,m2,m2,m2]
        #    [m1,m1,m1,m1,m1,m2,m2,m2,m2,m2]
        #    [m1,m1,m1,m1,m1,m2,m2,m2,m2,m2]
        #    [ 0, 0, 0, 0, 0,m2,m2,m2,m2,m2]
        #    [ 0, 0, 0, 0, 0,m2,m2,m2,m2,m2]]
        f_m1_m2 = Federated(self.sds,
            [fed1, fed2],
            [([0, 0], [dim, dim]), ([2, dim], [dim + 2, dim * 2])]
        ).compute()

        m1_p = np.concatenate((m1, np.zeros((2, dim))))
        m2_p = np.concatenate((np.zeros((2, dim)), m2))
        m1_m2 = np.concatenate((m1_p, m2_p), axis=1)
        res = np.allclose(f_m1_m2, m1_m2)
        self.assertTrue(res)

    # def test_6(self):
    #     # Note it overwrites the value in the field. not sum or anything else.
    #     #   [[m1,m1,m1,m1,m1, 0, 0, 0]
    #     #    [m1,m1,m1,m1,m1, 0, 0, 0]
    #     #    [m1,m1,m1,m2,m2,m2,m2,m2]
    #     #    [m1,m1,m1,m2,m2,m2,m2,m2]
    #     #    [m1,m1,m1,m2,m2,m2,m2,m2]
    #     #    [ 0, 0, 0,m2,m2,m2,m2,m2]
    #     #    [ 0, 0, 0,m2,m2,m2,m2,m2]]
    #     f_m1_m2 = Federated(self.sds,
    #         [fed1, fed2], [([0, 0], [dim, dim]), ([2, 3], [dim + 2, dim + 3])]
    #     ).compute()

    #     m1_m2 = np.zeros((dim + 2, dim + 3))
    #     m1_m2[0:dim, 0:dim] = m1
    #     m1_m2[2 : dim + 2, 3 : dim + 3] = m2

    #     res = np.allclose(f_m1_m2, m1_m2)
    #     self.assertTrue(res)

    # def test_7(self):
    #     #   [[m1,m1,m1,m1,m1, 0, 0, 0]
    #     #    [m1,m1,m1,m1,m1, 0, 0, 0]
    #     #    [m1,m1,m1,m2,m2,m2,m2,m2]
    #     #    [m1,m1,m1,m2,m2,m2,m2,m2]    +     1
    #     #    [m1,m1,m1,m2,m2,m2,m2,m2]
    #     #    [ 0, 0, 0,m2,m2,m2,m2,m2]
    #     #    [ 0, 0, 0,m2,m2,m2,m2,m2]]
    #     f_m1_m2 = Federated(self.sds,
    #         [fed1, fed2], [([0, 0], [dim, dim]), ([2, 3], [dim + 2, dim + 3])]
    #     )
    #     f_m1_m2 = (f_m1_m2 + 1).compute()
    #     m1_m2 = np.zeros((dim + 2, dim + 3))
    #     m1_m2[0:dim, 0:dim] = m1
    #     m1_m2[2 : dim + 2, 3 : dim + 3] = m2
    #     m1_m2 += 1
    #     res = np.allclose(f_m1_m2, m1_m2)
    #     if not res:
    #         print("Federated:")
    #         print(f_m1_m2)
    #         print("numpy:")
    #         print(m1_m2)
    #     self.assertTrue(res)

    def test_8(self):
        #   [[ 0, 0, 0, 0, 0, 0, 0, 0]
        #    [ 0, 0, 0, 0, 0, 0, 0, 0]
        #    [ 0, 0, 0,m1,m1,m1,m1,m1]
        #    [ 0, 0, 0,m1,m1,m1,m1,m1]
        #    [ 0, 0, 0,m1,m1,m1,m1,m1]
        #    [ 0, 0, 0,m1,m1,m1,m1,m1]
        #    [ 0, 0, 0,m1,m1,m1,m1,m1]]
        f_m1_m2 = Federated(self.sds, [fed1], [([2, 3], [dim + 2, dim + 3])])
        f_m1_m2 = (f_m1_m2).compute()
        m1_m2 = np.zeros((dim + 2, dim + 3))
        m1_m2[2: dim + 2, 3: dim + 3] = m1
        res = np.allclose(f_m1_m2, m1_m2)
        if not res:
            print("Federated:")
            print(f_m1_m2)
            print("numpy:")
            print(m1_m2)
        self.assertTrue(res)

    # def test_9(self):
    #     #   [[ 0, 0, 0, 0, 0, 0, 0, 0]
    #     #    [ 0, 0, 0, 0, 0, 0, 0, 0]
    #     #    [ 0, 0, 0,m1,m1,m1,m1,m1]
    #     #    [ 0, 0, 0,m1,m1,m1,m1,m1]    +     1
    #     #    [ 0, 0, 0,m1,m1,m1,m1,m1]
    #     #    [ 0, 0, 0,m1,m1,m1,m1,m1]
    #     #    [ 0, 0, 0,m1,m1,m1,m1,m1]]
    #     f_m1_m2 = Federated(self.sds, [fed1], [([2, 3], [dim + 2, dim + 3])])
    #     f_m1_m2 = (f_m1_m2 + 1).compute()

    #     m1_m2 = np.zeros((dim + 2, dim + 3))
    #     m1_m2[2 : dim + 2, 3 : dim + 3] = m1

    #     m1_m2 += 1
    #     res = np.allclose(f_m1_m2, m1_m2)

    #     if not res:
    #         print("Federated:")
    #         print(f_m1_m2)
    #         print("numpy:")
    #         print(m1_m2)
    #     self.assertTrue(res)

    # def test_10(self):
    #     #   [[m1,m1,m1,m1,m1, 0, 0, 0]
    #     #    [m1,m1,m1,m1,m1, 0, 0, 0]
    #     #    [m1,m1,m1,m1,m1, 0, 0, 0]
    #     #    [m1,m1,m1,m1,m1, 0, 0, 0]
    #     #    [m1,m1,m1,m1,m1, 0, 0, 0]
    #     #    [ 0, 0, 0, 0, 0, 0, 0, 0]
    #     #    [ 0, 0, 0, 0, 0, 0, 0, 0]]
    #     f_m1_m2 = Federated(self.sds, [fed1], [([0, 0], [dim + 2, dim + 3])])
    #     f_m1_m2 = (f_m1_m2).compute()

    #     m1_m2 = np.zeros((dim + 2, dim + 3))
    #     m1_m2[0:dim, 0:dim] = m1

    #     res = np.allclose(f_m1_m2, m1_m2)

    #     if not res:
    #         print("Federated:")
    #         print(f_m1_m2)
    #         print("numpy:")
    #         print(m1_m2)
    #     self.assertTrue(res)

    # def test_11(self):
    #     #   [[ 0, 0, 0, 0, 0, 0, 0, 0]
    #     #    [ 0,m1,m1,m1,m1,m1, 0, 0]
    #     #    [ 0,m1,m1,m1,m1,m1, 0, 0]
    #     #    [ 0,m1,m1,m1,m1,m1, 0, 0]
    #     #    [ 0,m1,m1,m1,m1,m1, 0, 0]
    #     #    [ 0,m1,m1,m1,m1,m1, 0, 0]
    #     #    [ 0, 0, 0, 0, 0, 0, 0, 0]]
    #     f_m1_m2 = Federated(self.sds, [fed1], [([1, 1], [dim + 2, dim + 3])])
    #     f_m1_m2 = (f_m1_m2).compute()

    #     m1_m2 = np.zeros((dim + 2, dim + 3))
    #     m1_m2[1 : dim + 1, 1 : dim + 1] = m1

    #     res = np.allclose(f_m1_m2, m1_m2)

    #     if not res:
    #         print("Federated:")
    #         print(f_m1_m2)
    #         print("numpy:")
    #         print(m1_m2)
    #     self.assertTrue(res)

if __name__ == "__main__":
    unittest.main(exit=False)
    shutil.rmtree(tempdir)


import os
import numpy as np

try:
    import flopy
except:
    msg = "Error. FloPy package is not available.\n"
    msg += "Try installing using the following command:\n"
    msg += " pip install flopy"
    raise Exception(msg)

from framework import testing_framework, running_on_CI
from simulation import Simulation

ex = ["gwf_noptc01", "gwf_noptc02", "gwf_noptc03"]
exdirs = []
for s in ex:
    exdirs.append(os.path.join("temp", s))

no_ptcrecords = ["FIRST", "ALL", None]

ddir = "data"

## run all examples on Travis
continuous_integration = [True for idx in range(len(exdirs))]

# set replace_exe to None to use default executable
# replace_exe = {'mf2005': 'mf2005devdbl'}
replace_exe = None

htol = [None for idx in range(len(exdirs))]

# static model data
# temporal discretization
nper = 1
tdis_rc = [(1.0, 1, 1.0)]

# spatial discretization data
nlay, nrow, ncol = 1, 1, 100
shape3d = (nlay, nrow, ncol)
size3d = nlay * nrow * ncol
delr, delc = 50.0, 1.0
top = 25.0
botm = 0.0
strt = 0.0

# hydraulic properties
hk = 50.0

# all cells are active and layer 1 is convertible
ib = 1

# solver options
nouter, ninner = 500, 300
hclose, rclose, relax = 1e-9, 1e-6, 1.0
newtonoptions = ""
imsla = "BICGSTAB"

# chd data
c6 = []
ccol = [0, ncol - 1]
hc = [20.0, 11.0]
for j, h in zip(ccol, hc):
    c6.append([(0, 0, j), h])
cd6 = {0: c6}
maxchd = len(cd6[0])

# recharge data
rech = {0: 0.001}

def build_model(idx, dir, no_ptcrecord):
    name = ex[idx]

    # build MODFLOW 6 files
    ws = dir
    sim = flopy.mf6.MFSimulation(
        sim_name=name, version="mf6", exe_name="mf6", sim_ws=ws
    )
    # create tdis package
    tdis = flopy.mf6.ModflowTdis(
        sim, time_units="DAYS", nper=nper, perioddata=tdis_rc
    )

    # create gwf model
    gwf = flopy.mf6.ModflowGwf(
        sim, modelname=name, newtonoptions=newtonoptions, save_flows=True
    )

    # create iterative model solution and register the gwf model with it
    ims = flopy.mf6.ModflowIms(
        sim,
        print_option="SUMMARY",
        no_ptcrecord=no_ptcrecord,
        outer_dvclose=hclose,
        outer_maximum=nouter,
        under_relaxation="NONE",
        inner_maximum=ninner,
        inner_dvclose=hclose,
        rcloserecord=rclose,
        linear_acceleration=imsla,
        scaling_method="NONE",
        reordering_method="NONE",
        relaxation_factor=relax,
    )
    sim.register_ims_package(ims, [gwf.name])

    dis = flopy.mf6.ModflowGwfdis(
        gwf,
        nlay=nlay,
        nrow=nrow,
        ncol=ncol,
        delr=delr,
        delc=delc,
        top=top,
        botm=botm,
    )

    # initial conditions
    ic = flopy.mf6.ModflowGwfic(gwf, strt=strt)

    # node property flow
    npf = flopy.mf6.ModflowGwfnpf(gwf, save_flows=False, icelltype=1, k=hk)

    # recharge
    rch = flopy.mf6.ModflowGwfrcha(gwf, readasarrays=True, recharge=rech)

    # chd files
    chd = flopy.mf6.modflow.mfgwfchd.ModflowGwfchd(
        gwf, maxbound=maxchd, stress_period_data=cd6, save_flows=False
    )

    # output control
    oc = flopy.mf6.ModflowGwfoc(
        gwf,
        budget_filerecord="{}.cbc".format(name),
        head_filerecord="{}.hds".format(name),
        headprintrecord=[("COLUMNS", 10, "WIDTH", 15, "DIGITS", 6, "GENERAL")],
        saverecord=[("HEAD", "ALL"), ("BUDGET", "ALL")],
        printrecord=[("HEAD", "LAST"), ("BUDGET", "ALL")],
    )

    return sim

# water table recharge problem
def get_model(idx, dir):
    sim = build_model(idx, dir, no_ptcrecords[idx])

    # build MODFLOW-6 without no_ptc option
    pth = os.path.join(dir, "mf6")
    mc = build_model(idx, pth, None)

    return sim, mc

def build_models():
    for idx, dir in enumerate(exdirs):
        sim, mc = get_model(idx, dir)
        sim.write_simulation()
        mc.write_simulation()
    return

# - No need to change any code below
def test_mf6model():
    # determine if running on Travis or GitHub actions
    is_CI = running_on_CI()
    r_exe = None
    if not is_CI:
        if replace_exe is not None:
            r_exe = replace_exe

    # initialize testing framework
    test = testing_framework()

    # build the models
    build_models()

    # run the test models
    for idx, dir in enumerate(exdirs):
        if is_CI and not continuous_integration[idx]:
            continue
        yield test.run_mf6, Simulation(dir, idxsim=idx)

    return

def main():
    # initialize testing framework
    test = testing_framework()

    # build the models
    build_models()

    # run the test models
    for idx, dir in enumerate(exdirs):
        sim = Simulation(dir, idxsim=idx)
        test.run_mf6(sim)

    return

# use python testmf6_csub_sub03.py --mf2005 mf2005devdbl
if __name__ == "__main__":
    # print message
    print("standalone run of {}".format(os.path.basename(__file__)))

    # run main routine
    main()