import os
import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt
#
from scipy.optimize import curve_fit
#
from regrowth_modeling_auxiliary_sctipt import gedi_icesat2_processor

# #####################################################################################
# parameter description
# #####################################################################################
# NOTE: all processing parameters are boolean type, i.e. True/False
# ------------------
# CALIBRATED_RH98 - if true, GEDI and ICESat-2 height values will be calibrated and calibrated heigths will be added
# as an extra column in the output pandas dataframes.
# ------------------
# SENS_ALG2 -  if true, only sensitivity from the algorithm setting group 2 will be used.
# (The auto-detected sensitivity will be ignored)
# ------------------
# GROUP2_ONLY - if true, only shots from the algorithm setting group 2 will be considered.
# ------------------
# SENS_SUBGROUPS_ONLY - if true, only subgroups with different sensitivity level (ALL, QS90, QS95, QS98, and QS99)
# will be returned. Keep this argument always set as True.
# #####################################################################################
# define different processing scenarios
# #####################################################################################
# my_ParSet_X = [CALIBRATED_RH98, SENS_ALG2, GROUP2_ONLY, SENS_SUBGROUPS_ONLY]
my_ParSet_1 = [True, True, False, True]
my_ParSet_2 = [True, False, False, True]
my_ParSet_3 = [True, False, True, True]
my_ParSet_4 = [False, False, False, True]
#
my_settings = [my_ParSet_1, my_ParSet_2, my_ParSet_3, my_ParSet_4]
my_ProcNames = ['Cal_SensA2', 'Cal', 'Cal_Group2', 'NotCal']

# #######################################################
# define the models
# #######################################################
# define my function for curve_fit
def my_fun1(xx, a, b):
    return a*xx + b

# specify max years to analyse
maxYears = 20
# #####################################################################################
# process the GEDI and ICESat-2 data
# #####################################################################################
# initiate empty dataframes:
df_r2 = pd.DataFrame({'Cal_SensA2': [], 'Cal': [], 'Cal_Group2': [], 'NotCal': []})
df_Intc = pd.DataFrame({'Cal_SensA2': [], 'Cal': [], 'Cal_Group2': [], 'NotCal': []})
df_Slope = pd.DataFrame({'Cal_SensA2': [], 'Cal': [], 'Cal_Group2': [], 'NotCal': []})
df_StdIntc = pd.DataFrame({'Cal_SensA2': [], 'Cal': [], 'Cal_Group2': [], 'NotCal': []})
df_StdSlope = pd.DataFrame({'Cal_SensA2': [], 'Cal': [], 'Cal_Group2': [], 'NotCal': []})
#
# loop trough each processing scenario
for my_setting, my_ProcName in zip(my_settings, my_ProcNames):
    [CALIBRATED_RH98, SENS_ALG2, GROUP2_ONLY, SENS_SUBGROUPS_ONLY] = my_setting
    # obtain the subgroups:
    [gedi_df_subgroups, ignore] = gedi_icesat2_processor(*my_setting)
    # specify the gedi subgroups' names:
    if SENS_SUBGROUPS_ONLY:
        gedi_df_names = ["ALL", "QS90", "QS95", "QS98", "QS99"]
    else:
        gedi_df_names = ["QS90-QPN", "QS90-QPD", "QS90-QPND",
                         "QS90-QCN", "QS90-QCD", "QS90-QCND",
                         "QS90-QPCN", "QS90-QPCD", "QS90"]

    # specify what RH98 value should be considered for regrowth rate:
    if CALIBRATED_RH98:
        height_col = 'rh_98_cal'
    else:
        height_col = 'rh_98'
    # #######################################################
    # Fit the linear model
    # #######################################################
    #
    my_r2 = np.full([len(gedi_df_names)], np.nan)
    my_Intercept = np.full([len(gedi_df_names)], np.nan)
    my_Slope = np.full([len(gedi_df_names)], np.nan)
    my_StdIntc = np.full([len(gedi_df_names)], np.nan)
    my_StdSlope = np.full([len(gedi_df_names)], np.nan)
    # ------------------------------------------------
    # loop trough each sensitivity subgroup (all, QS90, QS95, QS98, and QS99)
    for myIndx in np.arange(len(gedi_df_names)):
        # set the current gdf:
        my_gedi_df = gedi_df_subgroups[myIndx]
        # get the subgroup' data frames:
        my_gedi_df = gedi_df_subgroups[myIndx]
        # group the data by stand age:
        df_standAge_grouped_gedi = my_gedi_df.groupby('forestAge').quantile(0.5)
        # select only first
        df_standAge_grouped_gedi = df_standAge_grouped_gedi[0:maxYears]
        # prepare data for sklearn fit:
        xx = df_standAge_grouped_gedi.index.values.reshape((-1, 1))
        yy = df_standAge_grouped_gedi[height_col].values
        #weights1 = my_gedi_df.groupby('forestAge').count().id.values[0:17]
        weights1 = np.ones(yy.shape[0])
        # ------------------------------------------------
        # fit the weighted regression
        model1 = LinearRegression()
        model1.fit(xx, yy, weights1)
        # get the Standard errors from curve_fit:
        xx1 = df_standAge_grouped_gedi.index.values
        pout, pcov = curve_fit(f=my_fun1, xdata=xx1, ydata=yy)
        perr = np.sqrt(np.diag(pcov))
        #
        my_r2[myIndx] = model1.score(xx, yy, weights1)
        my_Intercept[myIndx] = model1.intercept_
        my_Slope[myIndx] = model1.coef_[0]
        my_StdIntc[myIndx] = perr[1]
        my_StdSlope[myIndx] = perr[0]

    df_Intc[my_ProcName] = my_Intercept
    df_Intc['name'] = gedi_df_names
    #
    df_r2[my_ProcName] = my_r2
    df_r2['name'] = gedi_df_names
    #
    df_Slope[my_ProcName] = my_Slope
    df_Slope['name'] = gedi_df_names
    #
    df_StdSlope[my_ProcName] = my_StdSlope
    df_StdSlope['name'] = gedi_df_names
    #
    df_StdIntc[my_ProcName] = my_StdIntc
    df_StdIntc['name'] = gedi_df_names

# #######################################################
# plot the fitting parameters as bar plots
# #######################################################
# set width of bars
barWidth = 0.14
# set heights of bars
bars1_Cal_a2 = df_r2['Cal_SensA2'].values
bars1_Cal = df_r2['Cal'].values
bars1_Cal_Gr2 = df_r2['Cal_Group2'].values
bars1_noCal = df_r2['NotCal'].values
#
bars2_Cal_a2 = df_Intc['Cal_SensA2'].values
bars2_Cal = df_Intc['Cal'].values
bars2_Cal_Gr2 = df_Intc['Cal_Group2'].values
bars2_noCal = df_Intc['NotCal'].values
#
bars3_Cal_a2 = df_Slope['Cal_SensA2'].values
bars3_Cal = df_Slope['Cal'].values
bars3_Cal_Gr2 = df_Slope['Cal_Group2'].values
bars3_noCal = df_Slope['NotCal'].values

#
# Set position of bar on X axis
r1 = np.arange(len(bars1_Cal_a2))
r2 = [x - barWidth for x in r1]
r3 = [x - barWidth for x in r2]
r4 = [x + barWidth for x in r1]

# Make the plot
fig1, ax1 = plt.subplots(1, 3, figsize=(12, 3.5))
ax1[0].bar(r1, bars1_Cal_a2, color='lightgreen', width=barWidth, edgecolor='black', label='QS9X-S2 Subgroups')
ax1[0].bar(r2, bars1_Cal, color='dimgrey', width=barWidth, edgecolor='black', label='QS9X Subgroups')
#ax1[0].bar(r3, bars1_Cal_Gr2, color='gold', width=barWidth, edgecolor='black', label='QS9X-A2 Subgroups')
ax1[0].bar(r4, bars1_noCal, color='tomato', width=barWidth, edgecolor='black', label='Not-calibrated Heights')
#
ax1[0].set_ylabel('R$^2$  ', rotation=0)
ax1[0].set_ylim([0.6, 1])
# -----------
ax1[1].bar(r1, bars2_Cal_a2, color='lightgreen',
           yerr=2*df_StdIntc['Cal_SensA2'].values, error_kw={'linewidth': 1, 'capsize': 1.5},
           width=barWidth, edgecolor='black', label='QS9X-S2 Subgroups')
ax1[1].bar(r2, bars2_Cal, color='dimgrey',
           yerr=2*df_StdIntc['Cal'].values, error_kw={'linewidth': 1, 'capsize': 1.5},
           width=barWidth, edgecolor='black', label='QS9X Subgroups')
# ax1[1].bar(r3, bars2_Cal_Gr2, color='gold',
#            yerr=2*df_StdIntc['Cal_Group2'].values, error_kw={'linewidth': 1, 'capsize': 1.5},
#            width=barWidth, edgecolor='black', label='QS9X-A2 Subgroups')
ax1[1].bar(r4, bars2_noCal, color='tomato',
           yerr=2*df_StdIntc['NotCal'].values, error_kw={'linewidth': 1, 'capsize': 1.5},
           width=barWidth, edgecolor='black', label='Not-calibrated Heights')
#
ax1[1].set_ylabel('Intercept [m]')
# -----------
ax1[2].bar(r1, bars3_Cal_a2, color='lightgreen',
           yerr=2*df_StdSlope['Cal_SensA2'].values, error_kw={'linewidth': 1, 'capsize': 1.5},
           width=barWidth, edgecolor='black', label='QS9X-S2 Subgroups')
ax1[2].bar(r2, bars3_Cal, color='dimgrey',
           yerr=2*df_StdSlope['Cal'].values, error_kw={'linewidth': 1, 'capsize': 1.5},
           width=barWidth, edgecolor='black', label='QS9X Subgroups')
# ax1[2].bar(r3, bars3_Cal_Gr2, color='gold',
#            yerr=2*df_StdSlope['Cal_Group2'].values, error_kw={'linewidth': 1, 'capsize': 1.5},
#            width=barWidth, edgecolor='black', label='QS9X-A2 Subgroups')
ax1[2].bar(r4, bars3_noCal, color='tomato',
           yerr=2*df_StdSlope['NotCal'].values, error_kw={'linewidth': 1, 'capsize': 1.5},
           width=barWidth, edgecolor='black', label='Not-calibrated Heights')
#
ax1[2].set_ylabel('Slope [m/year]')

# --------------
# Add xticks on the middle of the group bars
#ax1[0].set_xlabel('GEDI Subgroup')
my_ticks = ["ALL", "QS90", "QS95", "QS98", "QS99"]
ax1[0].set_xticks(r1 - barWidth / 2)
ax1[0].set_xticklabels(my_ticks)
#
ax1[1].set_xticks(r1 - barWidth / 2)
ax1[1].set_xticklabels(my_ticks)
#ax1[1].set_xlabel('GEDI Subgroup')
#
ax1[2].set_xticks(r1 - barWidth / 2)
ax1[2].set_xticklabels(my_ticks)
#
ax1[0].set_axisbelow(True)
ax1[0].yaxis.grid()
ax1[1].set_axisbelow(True)
ax1[1].yaxis.grid()
ax1[2].set_axisbelow(True)
ax1[2].yaxis.grid()
#
#ax1[0].text(-1.5, 0.7, r'(a)', fontsize=13)
#ax1[1].text(-1.5, 0, r'(b)', fontsize=13)
#ax1[2].text(-1.5, 0, r'(c)', fontsize=13)
#
ax1[1].legend(ncol=3, loc='upper center', bbox_to_anchor=(-0.6, 1.05, 2.26, .102), mode="expand", borderaxespad=0)
fig1.tight_layout()

# #######################
outFilePath = r'/mnt/raid/milutin/upScaling/Rondonia/GEDI/Rondonia_L2A_v002/output/directAnalysis/figures_regroth_20years/'
out_FigName = 'GEDI_modeling20years_differentProcessingStrategies.png'
plt.savefig(os.path.join(os.path.dirname(outFilePath), out_FigName), dpi=150)


import base64
from io import BytesIO

import pandas as pd
import matplotlib.pyplot as plt

def create_and_save_graphs(data_website_reason, data_website_duration):
    # Convert list of tuples to DataFrame
    df_reason = pd.DataFrame(data_website_reason, columns=['Website', 'Reason'])
    df_duration = pd.DataFrame(data_website_duration, columns=['Website', 'Duration'])

    # Merge the two dataframes on the 'Website' column
    merged_df = pd.merge(df_reason, df_duration, on='Website')

    # Convert seconds to minutes for plotting
    merged_df['Duration'] = merged_df['Duration'] / 60  # Conversion to minutes

    # Group by Reason and aggregate (summing duration in minutes)
    aggregated_data_reason = merged_df.groupby('Reason').agg({
        'Duration': 'sum'
    }).reset_index()

    # Plotting the bar graph for activity and duration
    fig1, ax1 = plt.subplots(figsize=(10,6))
    ax1.bar(aggregated_data_reason['Reason'], aggregated_data_reason['Duration'], color='skyblue')
    ax1.set_title('Duration of Different Activities')
    ax1.set_xlabel('Activity')
    ax1.set_ylabel('Duration (minutes)')
    ax1.set_xticks(aggregated_data_reason['Reason'])
    ax1.set_xticklabels(aggregated_data_reason['Reason'], rotation=45)
    ax1.grid(axis='y')

    # Group by Website and aggregate (summing duration in minutes)
    aggregated_data_website = df_duration.groupby('Website').agg({
        'Duration': 'sum'
    }).reset_index()

    # Plotting the pie chart for websites and time
    fig2, ax2 = plt.subplots(figsize=(10,6))
    ax2.pie(aggregated_data_website['Duration'], labels=aggregated_data_website['Website'], autopct='%1.1f%%', startangle=140, colors=plt.cm.Paired.colors)
    ax2.set_title('Proportion of Time Spent on Different Websites')
    ax2.axis('equal')

    # Save the bar graph as an image file
    buffer_bar = BytesIO()
    fig1.savefig(buffer_bar, format='png')
    buffer_bar.seek(0)
    bar_base64 = base64.b64encode(buffer_bar.read()).decode()

    buffer_pie = BytesIO()
    fig2.savefig(buffer_pie, format='png')
    buffer_pie.seek(0)
    pie_bas64 = base64.b64encode(buffer_pie.read()).decode()

    # Return the base64 of the saved graphs
    return bar_base64, pie_bas64


import time
from tracking import *
import os
import pandas as pd
from my_utils import *
import cv2
from setup_video import setup_video
import pickle
from model_setup import *

plt.style.use('fivethirtyeight')

np.random.seed(42)
fig = plt.figure()

if os.path.exists(outputVideoPath):
    original_output_path = outputVideoPath.split('.')
    version = 1
    while os.path.exists(outputVideoPath):
        outputVideoPath = f"{original_output_path[0]}_dashcam_{version}.avi"
        version += 1

# initialize the video stream, pointer to output video file, and
# frame dimensions
videoStream = cv2.VideoCapture(inputVideoPath)
video_width = int(videoStream.get(cv2.CAP_PROP_FRAME_WIDTH))
video_height = int(videoStream.get(cv2.CAP_PROP_FRAME_HEIGHT))

# Specifying coordinates for a default line
x1_line = 0
y1_line = video_height // 2
x2_line = video_width
y2_line = video_height // 2

# Initialization
all_detections = [{(0, 0): 0} for i in range(FRAMES_BEFORE_CURRENT)]

# previous_frame_detections = [spatial.KDTree([(0,0)])]*FRAMES_BEFORE_CURRENT # Initializing all trees
num_frames, vehicle_count = 0, 0
frame_num = 0
writer = create_video_writer(video_width, video_height, videoStream, outputVideoPath)
start_time = int(time.time())
# loop over frames from the video file stream
while True:
    batch_start = time.time()
    frames = []

    for i in range(BATCH_SIZE):
        for j in range(RUN_EVERY):
            (grabbed, frame) = videoStream.read()
            frame_num += 1
        if frame is None:
            break
        frames.append(frame)
    if not frames:
        break

    print(f"FRAME:{frame_num}")
    boxes_list, confidences_list, classes_list = get_model_output(model, frames)

    for boxes, confidences, classids, frame in zip(boxes_list, confidences_list, classes_list, frames):
        draw_dashcam_boxes(boxes, classids, confidences, frame)
        # idxs, \
        # vehicle_count, \
        # current_detections, \
        # current_positions = identify_vehicles(boxes,
        #                                       confidences,
        #                                       classids,
        #                                       frame,
        #                                       vehicle_count,
        #                                       all_detections)
        # all_detections.append(current_detections)

        #### FRAME TEXT ####
        # fps = BATCH_SIZE / (time.time() - batch_start)
        # frame_text = get_frame_text(lanes, lines, fps)
        # write_text(frame, frame_text)

        #### SHOW AND SAVE FRAME ####
        writer.write(frame)

    # df = pd.DataFrame([{label: center for center, label in prev.items()} for prev in all_detections])
    # df.to_pickle(f'stats_dfs/{outputVideoPath.split("/")[-1].split(".")[-2]}')
    # counts = get_counts(lanes, df)
    if show_video and frame_num > 2 * RUN_EVERY:
        # plt.xlim((0, frame_num / RUN_EVERY))
        # fig, img = create_plot_img(fig, counts)
        # cv2.imshow("Plot", img)
        cv2.imshow("Video Feed", frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
    if frame_num>120:
        break

# release the file pointers
# max_id = max([max(list(detections.values())) for detections in all_detections])
# lane_log = create_lane_log(lanes, max_id)
# lane_switches = get_switches(lane_log)
# pickle_dump(lane_switches, "switches")
# print(lane_switches)
# print("[INFO] cleaning up...")
# df.to_pickle('stats')
# pickle_dump(lanes, "lanes")
# pickle_dump(lines, "lines")
# writer.release()
# videoStream.release()


from django.utils.translation import ugettext_lazy as _
from admin_tools.dashboard import modules
from events.loaders import is_github_remote_enabled
from events.mailchimp_utils import is_mailchimp_configured
from itkpimail import settings

class ConfigurationChecker(modules.DashboardModule):
    deletable = False
    collapsible = False

    def is_empty(self):
        return False

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.title = _("Configuration check")
        self.template = 'my_blocks/checker.html'
        self.github_configured = False
        self.mailchimp_configured = False

    def init_with_context(self, context):
        super().init_with_context(context)
        request = context['request']
        self.github_configured = is_github_remote_enabled(request)
        self.github_configured_auth = (settings.GITHUB_API_TOKEN is not None)\
                                      and (settings.GITHUB_API_TOKEN.strip() != '')
        self.mailchimp_configured = is_mailchimp_configured(request)


'''
Import dependencies
'''
import os
import warnings
import time
import numpy as np
import random
from glob import glob
from PIL import Image
import matplotlib.pyplot as plt

from keras.initializers import TruncatedNormal
from keras.layers import Activation, BatchNormalization, Conv2D, Conv2DTranspose, Dense, Flatten, Input, LeakyReLU, Reshape
from keras.models import Model
from keras.optimizers import Adam
#
warnings.filterwarnings("ignore")

'''
Hyperparameters
'''
IMAGE_SIZE = 64
NOISE_SIZE = 100
LR_D = 0.00004
LR_G = 0.0004
BATCH_SIZE = 64
EPOCHS = 50
BETA1 = 0.5
WEIGHT_INIT_STDDEV = 0.02
EPSILON = 0.00005
SAMPLES_TO_SHOW = 5

'''
GAN
'''
def generator_model(z=(NOISE_SIZE,)):
    # Input Layer [100]
    input_layer=Input(z)
    # Fully connected layer [100 -> 4x4x1024]
    fully_connected = Dense(4*4*1024) (input_layer)
    fully_connected = Reshape((4, 4, 1024)) (fully_connected)
    # Default alpha of keras is 0.3;
    # Default alpha of tf is 0.2;
    # For consistency with original code, alpha defined as 0.2
    fully_connected = LeakyReLU(alpha=0.2) (fully_connected)

    # Convolutional Layer 1 [4x4x1024 -> 8x8x512]
    trans_conv1 = Conv2DTranspose(filters=512,
    			   			      kernel_size=[5,5],
	    					      strides=[2,2],
    						      padding="same",
    						      kernel_initializer=TruncatedNormal(stddev=WEIGHT_INIT_STDDEV)) (fully_connected)
    batch_trans_conv1 = BatchNormalization(epsilon=EPSILON) (trans_conv1)
    trans_conv1_out = LeakyReLU(alpha=0.2) (batch_trans_conv1)

    # Convolutional Layer 2 [8x8x512 -> 16x16x256]
    trans_conv2 = Conv2DTranspose(filters=256,
    						  	  kernel_size=[5,5],
	    					  	  strides=[2,2],
    						  	  padding="same",
    						  	  kernel_initializer=TruncatedNormal(stddev=WEIGHT_INIT_STDDEV)) (trans_conv1_out)
    batch_trans_conv2 = BatchNormalization(epsilon=EPSILON) (trans_conv2)
    trans_conv2_out = LeakyReLU(alpha=0.2) (batch_trans_conv2)

    # Convolutional Layer 3 [16x16x256 -> 32x32x128]
    trans_conv3 = Conv2DTranspose(filters=128,
    						  	  kernel_size=[5,5],
	    					  	  strides=[2,2],
    						  	  padding="same",
    						  	  kernel_initializer=TruncatedNormal(stddev=WEIGHT_INIT_STDDEV)) (trans_conv2_out)
    batch_trans_conv3 = BatchNormalization(epsilon=EPSILON) (trans_conv3)
    trans_conv3_out = LeakyReLU(alpha=0.2) (batch_trans_conv3)

    # Convolutional Layer 4 [32x32x128 -> 64x64x64]
    trans_conv4 = Conv2DTranspose(filters=64,
    						  	  kernel_size=[5,5],
	    					  	  strides=[2,2],
    						  	  padding="same",
    						  	  kernel_initializer=TruncatedNormal(stddev=WEIGHT_INIT_STDDEV)) (trans_conv3_out)
    batch_trans_conv4 = BatchNormalization(epsilon=EPSILON) (trans_conv4)
    trans_conv4_out = LeakyReLU(alpha=0.2) (batch_trans_conv4)

    # Output layer [64x64x64 -> 64x64x3]
    logits = Conv2DTranspose(filters=3,
    						 kernel_size=[5,5],
	    					 strides=[1,1],
    						 padding="same",
    						 kernel_initializer=TruncatedNormal(stddev=WEIGHT_INIT_STDDEV)) (trans_conv4_out)
    out = Activation('tanh') (logits)

    #
    model = Model(inputs=input_layer, outputs=out)
    model.summary()
    return model

def discriminator_model(input_shape=(IMAGE_SIZE, IMAGE_SIZE, 3)):
    # Input Layer [64x64x3]
    input_layer = Input(input_shape)

	# Convolutional Layer 1 [64x64x3 -> 32x32x64]
    conv1 = Conv2D(filters=64,
				   kernel_size=[5,5],
				   strides=[2,2],
				   padding="same",
    			   kernel_initializer=TruncatedNormal(stddev=WEIGHT_INIT_STDDEV)) (input_layer)
    batch_norm1 = BatchNormalization(epsilon=EPSILON) (conv1)
    conv1_out = LeakyReLU(alpha=0.2) (batch_norm1)

	# Convolutional Layer 2 [32x32x64 -> 16x16x128]
    conv2 = Conv2D(filters=128,
				   kernel_size=[5,5],
				   strides=[2,2],
				   padding="same",
    			   kernel_initializer=TruncatedNormal(stddev=WEIGHT_INIT_STDDEV)) (conv1_out)
    batch_norm2 = BatchNormalization(epsilon=EPSILON) (conv2)
    conv2_out = LeakyReLU(alpha=0.2) (batch_norm2)

	# Convolutional Layer 3 [16x16x128 -> 8x8x256]
    conv3 = Conv2D(filters=256,
				   kernel_size=[5,5],
				   strides=[2,2],
				   padding="same",
    			   kernel_initializer=TruncatedNormal(stddev=WEIGHT_INIT_STDDEV)) (conv2_out)
    batch_norm3 = BatchNormalization(epsilon=EPSILON) (conv3)
    conv3_out = LeakyReLU(alpha=0.2) (batch_norm3)

	# Convolutional Layer 4 [8x8x256 -> 8x8x512]
    conv4 = Conv2D(filters=512,
				   kernel_size=[5,5],
				   strides=[1,1],
				   padding="same",
    			   kernel_initializer=TruncatedNormal(stddev=WEIGHT_INIT_STDDEV)) (conv3_out)
    batch_norm4 = BatchNormalization(epsilon=EPSILON) (conv4)
    conv4_out = LeakyReLU(alpha=0.2) (batch_norm4)

	# Convolutional Layer 5 [8x8x512 -> 4x4x1024]
    conv5 = Conv2D(filters=1024,
				   kernel_size=[5,5],
				   strides=[2,2],
				   padding="same",
    			   kernel_initializer=TruncatedNormal(stddev=WEIGHT_INIT_STDDEV)) (conv4_out)
    batch_norm5 = BatchNormalization(epsilon=EPSILON) (conv5)
    conv5_out = LeakyReLU(alpha=0.2) (batch_norm5)

    # Flattened Convolutional Filter Output Layer [4x4x1024 -> 16384] ###?#?#?#?#??#
    flatten = Flatten() (conv5_out)

    # Fully connected layer [16384 -> 1]
    out = Dense(1, activation='sigmoid') (flatten)

    # ??????
    model = Model(inputs=input_layer, outputs=out)
    model.summary()
    return model

# Discriminator
discriminator = discriminator_model((IMAGE_SIZE, IMAGE_SIZE, 3))
discriminator.compile(loss='binary_crossentropy',optimizer=Adam(lr=LR_D, beta_1=BETA1),metrics=['accuracy'])
discriminator.trainable = False

# Generator
generator = generator_model((NOISE_SIZE,))

# GAN
gan_input = Input(shape=(NOISE_SIZE,))
x = generator(gan_input)
gan_out = discriminator(x)
gan = Model(gan_input, gan_out)
gan.summary()

gan.compile(loss='binary_crossentropy',optimizer=Adam(lr=LR_G, beta_1=BETA1))

'''
PLOT SAMPLE IMAGES
'''
def show_samples(sample_images, name, epoch):
    figure, axes = plt.subplots(1, len(sample_images), figsize = (IMAGE_SIZE, IMAGE_SIZE))
    for index, axis in enumerate(axes):
        axis.axis('off')
        image_array = sample_images[index]
        axis.imshow(image_array)
        image = Image.fromarray(image_array)
        image.save(name+"_"+str(epoch)+"_"+str(index)+".png")
    plt.savefig(name+"_"+str(epoch)+".png", bbox_inches='tight', pad_inches=0)
    plt.show()
    plt.close()

'''

'''
def test(input_z, epoch, OUTPUT_DIR):
    samples = generator.predict(input_z[:SAMPLES_TO_SHOW])
    sample_images = [((sample + 1.0) * 127.5).astype(np.uint8) for sample in samples]
    show_samples(sample_images, OUTPUT_DIR + "samples", epoch)

def summarize_epoch(d_losses, g_losses , data_shape, epoch, duration, input_z, OUTPUT_DIR):
    minibatch_size = int(data_shape[0]//BATCH_SIZE)
    print("Epoch {}/{}".format(epoch, EPOCHS),
          "\nDuration: {:.5f}".format(duration),
          "\nD Loss: {:.5f}".format(np.mean(d_losses[-minibatch_size:])),
          "\nG Loss: {:.5f}".format(np.mean(g_losses[-minibatch_size:])))
    fig, ax = plt.subplots()
    plt.plot(d_losses, label='Discriminator', alpha=0.6)
    plt.plot(g_losses, label='Generator', alpha=0.6)
    plt.title("Losses")
    plt.legend()
    plt.savefig(OUTPUT_DIR + "losses_" + str(epoch) + ".png")
    plt.show()
    plt.close()
    test(input_z, epoch, OUTPUT_DIR)

def get_batches(data):
    batches = []
    for i in range(int(data.shape[0]//BATCH_SIZE)):
        batch = data[i * BATCH_SIZE:(i + 1) * BATCH_SIZE]
        augmented_images = []
        for img in batch:
            image = Image.fromarray(img)
            if random.choice([True, False]):
                image = image.transpose(Image.FLIP_LEFT_RIGHT)
            augmented_images.append(np.asarray(image))
        batch = np.asarray(augmented_images)
        normalized_batch = (batch / 127.5) - 1.0
        batches.append(normalized_batch)
    return np.array(batches)

# Training
def train():
    '''
LOAD DATA
'''
    INPUT_DATA_DIR = "./simpsons-faces/cropped/" # Path to the folder with input images. For more info check simspons_dataset.txt
    OUTPUT_DIR = './model/'
    if not os.path.exists(OUTPUT_DIR):
        os.makedirs(OUTPUT_DIR)

    # Training
    # Import Data
    exclude_img = ["9746","9731","9717","9684","9637","9641","9642","9584","9541","9535",
    "9250","9251","9252","9043","8593","8584","8052","8051","8008","7957",
    "7958""7761","7762","9510","9307","4848","4791","4785","4465","2709",
    "7724","7715","7309","7064","7011","6961","6962","6963","6960","6949",
    "6662","6496","6409","6411","6406","6407","6170","6171","6172","5617",
    "4363","4232","4086","4047","3894","3889","3493","3393","3362","2780",
    "2710","2707","2708","2711","2712","2309","2056","1943","1760","1743",
    "1702","1281","1272","772","736","737","691","684","314","242","191"]
    exclude_img = [s + ".png" for s in exclude_img]
    input_images = np.asarray([np.asarray(Image.open(file).resize((IMAGE_SIZE, IMAGE_SIZE))) for file in glob(INPUT_DATA_DIR + '*') if file not in exclude_img])
    print ("Image Samples\nInput: " + str(input_images.shape))
    np.random.shuffle(input_images)

    print("Training Starts!")
    # Load weights from last completed epoch if they exists
    discriminator_weights = glob('discriminator_simpson_weights.h5_*')
    gan_weights = glob('generator_simpson_weights.h5_*')
    if discriminator_weights:
        discriminator_weights.sort()
        discriminator.load_model(discriminator_weights[-1])
    if gan_weights:
        gan_weights.sort()
        gan.load_model(gan_weights[-1])

    d_losses = []
    g_losses = []
    cum_d_loss = 0
    cum_g_loss = 0

    #EPOCHS = 50
    for epoch in range(EPOCHS):
        epoch += 1
        start_time = time.time()
        for batch_images in get_batches(input_images):
            noise_data = np.random.normal(0, 1, size=(BATCH_SIZE, NOISE_SIZE))

            # We use same labels for generated images as in the real training batch
            generated_images = generator.predict(noise_data)

            # Prepare real target labels
            noise_prop = 0.05 # Randomly flip 5% of targets
            real_labels = np.zeros((BATCH_SIZE, 1)) + np.random.uniform(low=0.0, high=0.1, size=(BATCH_SIZE, 1))
            flipped_idx = np.random.choice(np.arange(len(real_labels)), size=int(noise_prop*len(real_labels)))
            real_labels[flipped_idx] = 1 - real_labels[flipped_idx]

            # Train discriminator on real data
            d_loss_real = discriminator.train_on_batch(batch_images, real_labels)

            # Prepare labels for generated data
            fake_labels = np.ones((BATCH_SIZE, 1)) - np.random.uniform(low=0.0, high=0.1, size=(BATCH_SIZE, 1))
            flipped_idx = np.random.choice(np.arange(len(fake_labels)), size=int(noise_prop*len(fake_labels)))
            fake_labels[flipped_idx] = 1 - fake_labels[flipped_idx]

            # Train discriminator on generated data
            d_loss_fake = discriminator.train_on_batch(generated_images, fake_labels)
            d_loss = 0.5 * np.add(d_loss_real, d_loss_fake)
            cum_d_loss += d_loss
            d_losses.append(d_loss[0])

            # Train generator
            noise_data = np.random.normal(0, 1, size=(BATCH_SIZE, NOISE_SIZE))
            g_loss = gan.train_on_batch(noise_data, np.zeros((BATCH_SIZE, 1)))
            cum_g_loss += g_loss
            g_losses.append(g_loss)

        if epoch > 0 and epoch % 5 == 0 :
            print("saving model")
            discriminator.save_weights(OUTPUT_DIR + "discriminator_simpson_weights.h5_" + str(epoch))
            gan.save_weights(OUTPUT_DIR + "gan_simpson_model.h5_" + str(epoch))
            summarize_epoch(d_losses, g_losses, input_images.shape, epoch, time.time()-start_time, noise_data, OUTPUT_DIR)

train()