# -*- encoding: utf8 -*-

import pandas as pd
from sqlalchemy import create_engine
import numpy as np

pd.options.mode.chained_assignment = None

engine_audit = create_engine('mysql+mysqldb://yulong:ohVU1nPITynSZwR2@172.16.3.201:3306/new_transaction?charset=utf8',
                             echo=True)

chunksize = 5000
sql1 = """
select DISTINCT plan_id 
from baitiao_repay_plan_repay_record_ref 
where ref_amount > 0 
and created_at>='%s'
 AND repay_channel in(3,5,7,13,22,23,24)
"""
sql1_1 = """select id, plan_id, ref_amount from baitiao_repay_plan_repay_record_ref where plan_id in %s and ref_amount > 0"""
sql2 = """
select id plan_id, principle, interest, service_fee, should_punish, repay_status,red_package 
from new_transaction.user_bt_repayment_plan 
where id in %s """
sql_update = """
    update baitiao_repay_plan_repay_record_ref set principle = %s, interest = %s, service_fee = %s, punish = %s, collection_relief = %s, remain_income = %s,
    mitigate_principle = %s, mitigate_interest = %s, mitigate_service_fee = %s, mitigate_collection_relief = %s,red_package_fee=%s,mitigate_all=%s  where id = %s
    """
columns = ['id', 'plan_id', 'ref_amount', 'principle', 'interest', 'service_fee', 'should_punish', 'repay_status',
           're_principle', 're_interest', 're_service_fee', 're_punish', 're_relief', 're_remain',
           'mitigate_prin', 'mitigate_inst', 'mitigate_fee', 'mitigate_punish', 'red_package_fee', 'mitigate_all']
up_columns = ['re_principle', 're_interest', 're_service_fee', 're_punish', 're_relief', 're_remain',
              'mitigate_prin', 'mitigate_inst', 'mitigate_fee', 'mitigate_punish', 'red_package_fee', 'mitigate_all',
              'id']


def split_money_3(df):
    df['re_principle'] = 0
    df['re_interest'] = 0
    df['re_service_fee'] = 0
    df['re_punish'] = 0
    df['re_relief'] = 0
    df['re_remain'] = 0
    df['mitigate_prin'] = 0
    df['mitigate_inst'] = 0
    df['mitigate_fee'] = 0
    df['mitigate_punish'] = 0
    df['red_package_fee'] = 0
    df['mitigate_all'] = 0

    left_prin = round(df['principle'].values[0], 2)
    left_inst = round(df['interest'].values[0], 2)
    left_fee = round(df['service_fee'].values[0], 2)
    left_punish = round(df['should_punish'].values[0], 2)
    # 红包
    left_red_packet = round(df['red_package'].values[0], 2)
    for index, row in df.iterrows():
        ref_amount = round(row['ref_amount'], 2)
        # 分配本金
        if left_prin > 0:
            if ref_amount <= left_prin:
                df.loc[index, 're_principle'] = ref_amount
                left_prin -= ref_amount
                continue
            else:
                df.loc[index, 're_principle'] = left_prin
                ref_amount -= left_prin
                left_prin = 0
        # 分配利息
        if left_inst > 0:
            if ref_amount <= left_inst:
                df.loc[index, 're_interest'] = ref_amount
                left_inst -= ref_amount
                continue
            else:
                df.loc[index, 're_interest'] = left_inst
                ref_amount -= left_inst
                left_inst = 0
        # 分配服务费
        if left_fee > 0:
            if ref_amount <= left_fee:
                df.loc[index, 're_service_fee'] = ref_amount
                left_fee -= ref_amount
                continue
            else:
                df.loc[index, 're_service_fee'] = left_fee
                ref_amount -= left_fee
                left_fee = 0
        # 分配罚息
        if left_punish > 0:
            if ref_amount <= left_punish:
                df.loc[index, 're_punish'] = ref_amount
                left_punish -= ref_amount
                continue
            else:
                df.loc[index, 're_punish'] = left_punish
                ref_amount -= left_punish
                left_punish = 0
        # 多出来的放到剩余金额里
        df.loc[index, 're_remain'] = ref_amount
    # 修改最后一条记录的金额
    df.loc[index, 're_principle'] = round(df['principle'].values[0], 2) - round(sum(df['re_principle'].values[:-1]), 2)
    df.loc[index, 're_interest'] = round(df['interest'].values[0], 2) - round(sum(df['re_interest'].values[:-1]), 2)
    df.loc[index, 're_service_fee'] = round(df['service_fee'].values[0], 2) - round(
        sum(df['re_service_fee'].values[:-1]), 2)
    df.loc[index, 're_punish'] = round(
        round(df['should_punish'].values[0], 2) - round(sum(df['re_punish'].values[:-1]), 2), 2)
    sum_amount = df.loc[index, 're_principle'] + df.loc[index, 're_interest'] + df.loc[index, 're_service_fee'] + \
                 df.loc[index, 're_punish']

    sum_prin = sum(df['re_principle'])
    sum_inst = sum(df['re_interest'])
    sum_fee = sum(df['re_service_fee'])
    sum_punish = sum(df['re_punish'])
    sum_total = sum(df['ref_amount'])
    if sum_amount > df.loc[index, 'ref_amount']:  # 应还 > 实还，需要进行减免
        df.loc[index, 're_relief'] = round(sum_amount - df.loc[index, 'ref_amount'], 2)
        df.loc[index, 'mitigate_prin'] = round(max(sum_prin - sum_total, 0), 2)
        df.loc[index, 'mitigate_inst'] = round(max(sum_prin + sum_inst - df.loc[index, 'mitigate_prin'] - sum_total, 0),
                                               2)
        df.loc[index, 'mitigate_fee'] = round(max(
            sum_prin + sum_inst + sum_fee - df.loc[index, 'mitigate_prin'] - df.loc[index, 'mitigate_inst'] - sum_total,
            0), 2)
        df.loc[index, 'mitigate_punish'] = round(max(
            sum_prin + sum_inst + sum_fee + sum_punish - df.loc[index, 'mitigate_prin'] - df.loc[
                index, 'mitigate_inst'] - df.loc[index, 'mitigate_fee'] - sum_total, 0), 2)
        df.loc[index, 're_remain'] = 0
        # 红包,必须是已还款的
        if (df.loc[index, 're_relief'] >= left_red_packet) & (df.loc[index, 'repay_status'] == 3):
            df.loc[index, 'red_package_fee'] = left_red_packet
        else:
            # TODO 输出到log 文件中
            list_out.append(df.loc[index])
        df.loc[index, 'mitigate_all'] = df.loc[index, 're_relief'] - df.loc[index, 'red_package_fee']
    elif sum_amount < df.loc[index, 'ref_amount']:
        df.loc[index, 're_relief'] = 0
        df.loc[index, 're_remain'] = round(df.loc[index, 'ref_amount'] - sum_amount, 2)
    else:
        df.loc[index, 're_relief'] = 0
        df.loc[index, 're_remain'] = 0
    return df


def split_money_0(df):
    df['re_principle'] = 0
    df['re_interest'] = 0
    df['re_service_fee'] = 0
    df['re_punish'] = 0
    df['re_relief'] = 0
    df['re_remain'] = 0

    left_prin = round(df['principle'].values[0], 2)
    left_inst = round(df['interest'].values[0], 2)
    left_fee = round(df['service_fee'].values[0], 2)
    left_punish = round(df['should_punish'].values[0], 2)
    for index, row in df.iterrows():
        ref_amount = round(row['ref_amount'], 2)
        # 分配本金
        if left_prin > 0:
            if ref_amount <= left_prin:
                df.loc[index, 're_principle'] = ref_amount
                left_prin -= ref_amount
                continue
            else:
                df.loc[index, 're_principle'] = left_prin
                ref_amount -= left_prin
                left_prin = 0
        # 分配利息
        if left_inst > 0:
            if ref_amount <= left_inst:
                df.loc[index, 're_interest'] = ref_amount
                left_inst -= ref_amount
                continue
            else:
                df.loc[index, 're_interest'] = left_inst
                ref_amount -= left_inst
                left_inst = 0
        # 分配服务费
        if left_fee > 0:
            if ref_amount <= left_fee:
                df.loc[index, 're_service_fee'] = ref_amount
                left_fee -= ref_amount
                continue
            else:
                df.loc[index, 're_service_fee'] = left_fee
                ref_amount -= left_fee
                left_fee = 0
        # 分配罚息
        if left_punish > 0:
            if ref_amount <= left_punish:
                df.loc[index, 're_punish'] = ref_amount
                left_punish -= ref_amount
                continue
            else:
                df.loc[index, 're_punish'] = left_punish
                ref_amount -= left_punish
                left_punish = 0
        # 多出来的放到剩余金额里
        df.loc[index, 're_remain'] = ref_amount
    return df


start_time = '2017-10-04'
df1_gen = pd.read_sql(sql=sql1 % start_time, con=engine_audit, chunksize=chunksize)
i = 0
# 输出有问题的红包
list_out = []
for df1 in df1_gen:
    print '--------------------> %s <----------------------------' % (i,)
    i += 1
    plan_ids = df1['plan_id'].astype(str).tolist()
    df1 = pd.read_sql(sql=sql1_1 % (str(tuple(plan_ids)).replace(',)', ')'),), con=engine_audit)
    df = pd.read_sql(sql=sql2 % (str(tuple(plan_ids)).replace(',)', ')'),), con=engine_audit)

    df.fillna(0, inplace=True)
    df['repay_status'] = df['repay_status'].astype(int)
    df = pd.merge(df1, df, on='plan_id', how='left')
    df_cnt = df.groupby('plan_id')['id'].agg(['count']).reset_index()
    df = pd.merge(df, df_cnt, on='plan_id')

    df['re_principle'] = 0
    df['re_interest'] = 0
    df['re_service_fee'] = 0
    df['re_punish'] = 0
    df['re_relief'] = 0
    df['re_remain'] = 0
    df['mitigate_prin'] = 0
    df['mitigate_inst'] = 0
    df['mitigate_fee'] = 0
    df['mitigate_punish'] = 0
    df['red_package_fee'] = 0
    df['mitigate_all'] = 0

    # count=1的可以直接拆
    # count>1的需要合并拆
    df_3_1 = df.ix[(df['repay_status'] == 3) & (df['count'] == 1)]
    df_3_2 = df.ix[(df['repay_status'] == 3) & (df['count'] > 1)]
    df_0_1 = df.ix[(df['repay_status'] != 3) & (df['count'] == 1)]
    df_0_2 = df.ix[(df['repay_status'] != 3) & (df['count'] > 1)]

    if len(df_3_1) > 0:
        df_3_1['re_principle'] = df_3_1['principle']
        df_3_1['re_interest'] = df_3_1['interest']
        df_3_1['re_service_fee'] = df_3_1['service_fee']
        df_3_1['re_punish'] = df_3_1['should_punish']
        df_3_1['re_punish'] = df_3_1['should_punish'].apply(lambda x: round(x, 2))
        df_3_1['diff'] = df_3_1['ref_amount'] - (
            df_3_1['re_principle'] + df_3_1['re_interest'] + df_3_1['re_service_fee'] + df_3_1['re_punish'])
        df_3_1['re_relief'] = df_3_1['diff'].apply(lambda x: round(-x, 2) if x < 0 else 0)
        df_3_1['re_remain'] = df_3_1['diff'].apply(lambda x: round(x, 2) if x > 0 else 0)
        df_3_1['mitigate_prin'] = df_3_1.apply(lambda x: max(x['re_principle'] - x['ref_amount'], 0), axis=1)
        df_3_1['mitigate_inst'] = df_3_1.apply(
            lambda x: round(max(x['re_principle'] + x['re_interest'] - x['mitigate_prin'] - x['ref_amount'], 0), 2),
            axis=1)
        df_3_1['mitigate_fee'] = df_3_1.apply(lambda x: max(
            x['re_principle'] + x['re_interest'] + x['re_service_fee'] - x['mitigate_prin'] - x['mitigate_inst'] - x[
                'ref_amount'], 0), axis=1)
        df_3_1['mitigate_punish'] = df_3_1.apply(lambda x: max(
            x['re_principle'] + x['re_interest'] + x['re_service_fee'] + x['re_punish'] - x['mitigate_prin'] - x[
                'mitigate_inst'] - x['mitigate_fee'] - x['ref_amount'], 0), axis=1)
        df_3_1.ix[((df_3_1.re_relief >= df_3_1.red_package) & (df_3_1.repay_status == 3)), 'red_package_fee'] = \
        df_3_1.ix[((df_3_1.re_relief >= df_3_1.red_package) & (df_3_1.repay_status == 3))]['red_package']
        df_3_1['mitigate_all'] = df_3_1['re_relief'] - df_3_1['red_package_fee']
        # TODO 输出 》减免的红包
        list_out.append(df_3_1.ix[df_3_1.re_relief < df_3_1.red_package])
    if len(df_0_1) > 0:
        df_0_1['re_principle'] = df_0_1.apply(
            lambda x: round(x['principle'], 2) if x['ref_amount'] > x['principle'] else round(x['ref_amount'], 2),
            axis=1)
        df_0_1['ref_amount'] = np.round(df_0_1['ref_amount'] - df_0_1['principle'], 2)

        df_0_1.ix[df_0_1['ref_amount'] > 0, 're_interest'] = df_0_1.ix[df_0_1['ref_amount'] > 0].apply(
            lambda x: round(x['interest'], 2) if x['ref_amount'] > x['interest'] else x['ref_amount'],
            axis=1)
        df_0_1['ref_amount'] = np.round(df_0_1['ref_amount'] - df_0_1['interest'], 2)

        df_0_1.ix[df_0_1['ref_amount'] > 0, 're_service_fee'] = df_0_1.ix[df_0_1['ref_amount'] > 0].apply(
            lambda x: round(x['service_fee'], 2) if x['ref_amount'] > x['service_fee'] else x['ref_amount'],
            axis=1)
        df_0_1['ref_amount'] = np.round(df_0_1['ref_amount'] - df_0_1['service_fee'], 2)
        df_0_1['should_punish'] = df_0_1['should_punish'].apply(lambda x: round(x, 2))
        df_0_1.ix[df_0_1['ref_amount'] > 0, 're_punish'] = df_0_1.ix[df_0_1['ref_amount'] > 0].apply(
            lambda x: round(x['should_punish'], 2) if x['ref_amount'] > x['should_punish'] else x['ref_amount'],
            axis=1)
        df_0_1['ref_amount'] = np.round(df_0_1['ref_amount'] - df_0_1['should_punish'], 2)

        df_0_1.ix[df_0_1['ref_amount'] > 0, 're_remain'] = df_0_1.ix[df_0_1['ref_amount'] > 0, 'ref_amount']

    print '-------------------> len(df): %s <------------------------------------' % (len(df),)
    print '-------------------> len(df_3_1): %s, len(df_3_2): %s <------------------------------------' % (
    len(df_3_1), len(df_3_2))
    print '-------------------> len(df_0_1): %s, len(df_0_2): %s <------------------------------------' % (
    len(df_0_1), len(df_0_2))

    result_3 = pd.DataFrame(columns=columns)
    str1 = df_3_2.plan_id.unique()
    for i in str1:
        df = df_3_2.ix[df_3_2.plan_id == i]
        b = split_money_3(df)
        result_3 = result_3.append(b)
    # result_3.to_csv("df_3_2.csv", index=False, columns=columns)
    # df_3_1.to_csv("df_3_1.csv", index=False, columns=columns)

    print '-------------------> split 3, len(result_3): %s <------------------------------------' % (
        len(result_3),)
    result_0 = pd.DataFrame(columns=columns)
    str1 = df_0_2.plan_id.unique()
    for i in str1:
        df = df_0_2.ix[df_0_2.plan_id == i]
        b = split_money_0(df)
        result_0 = result_0.append(b)
    # result_0.to_csv("df_0_2.csv", index=False, columns=columns)
    # df_0_1.to_csv("df_0_1.csv", index=False, columns=columns)
    print '-------------------> split 0, len(result_0): %s <------------------------------------' % (
        len(result_0),)

    df_list = result_3[up_columns].values.tolist()
    df_list.extend(df_3_1[up_columns].values.tolist())
    df_list.extend(result_0[up_columns].values.tolist())
    df_list.extend(df_0_1[up_columns].values.tolist())
    print str(df_list)
    print sql_update
    engine_audit.execute(sql_update, df_list)
    print '-------------------> update len(df_list): %s <------------------------------------' % (len(df_list),)

print '----------------------------------> done ! <----------------------------------------------------'
