import csv
import functions
import Sector_Selection
import datetime
import os
import matplotlib.pyplot as plt

# Variables
initial_amount = 1000000  # 1M

# Initialization
market = "HSCI"
oneyear = 252  # 1 year = 252 trading days
numberofday = 0
timestamp = []
date_ranking = []
date_marketcap = []
date_index = []
quarter_list = []
cumulative_list = []
store_index = []  # when quarter changes
price_index = []
daily_return = []
portfolio_list = [initial_amount]
# Functions
replace = functions.replace_na
average = functions.findaverage
SD = functions.findSD
maxconsecutive = functions.MaxConsecutive
winrate = functions.WinRate
sharpe = functions.Sharpe
market_cap = functions.Check_number(Sector_Selection.market_cap)

# Store the date of market cap
for i in range(len(market_cap)):
    date_marketcap.append(market_cap[i][0])

# Set the directory
directory = "Sector Selection Result/"
if not os.path.exists(directory):
    os.makedirs(directory)

# Import market index and divide them by the first element
with open(market+'.csv') as f:
    reader = csv.reader(f)
    index = [rows for rows in reader]
    for i in range(len(index)):
        date_index.append(index[i][0])
        price_index.append(float(index[i][1]))

# Import Ranking generated by Sector_Selection
with open(directory + 'Ranking [1 and 0].csv') as f:
    reader = csv.reader(f)
    ranking = [rows for rows in reader]
    ranking = ranking[1:]
    for i in range(len(ranking)):
        date_ranking.append(ranking[i][0])

# Import prices of each stock
with open('prices.csv') as g:
    reader = csv.reader(g)
    price = [rows for rows in reader]
    price = replace(price)
    price = price[1:]
    for i in range(len(price)):
        timestamp.append(price[i][0])

print("\n")
print(market, "BACK TEST")
# Assign Time Stamp into correct quarter
ticker = Sector_Selection.ticker
sector = Sector_Selection.sector

# Initialize matrix full of zeros
resultMatrix = [[0 for x in range(len(ticker)+5)] for y in range(len(timestamp)+1)]

# Change the first row of result matrix to show headers
firstrow = resultMatrix[0]
secondrow = resultMatrix[1]
thirdrow = resultMatrix[2]
for i in range(len(firstrow)):
    if i == 0:
        firstrow[i] = ""
    elif i < len(ticker)+1:
        firstrow[i] = ticker[i-1]
    elif i == len(ticker)+1:
        firstrow[i] = "Daily Return"
    elif i == len(ticker)+2:
        firstrow[i] = "Cumulative"
    elif i == len(ticker)+3:
        firstrow[i] = "Stock Count"
    elif i == len(ticker)+4:
        firstrow[i] = "Portfolio"
        secondrow[i] = "Start From Below"
        thirdrow[i] = initial_amount

# Change the first column of result matrix to show time index
for i in range(len(resultMatrix)):
    if i == 0:
        resultMatrix[i][0] = ""
    else:
        resultMatrix[i][0] = timestamp[i-1]

# Loop through time index
for index in range(len(timestamp)):
    stockcount = 0
    answer = 0  # search for index for ranking
    answer_marketcap = 0
    market_cap_total = 0
    market_cap_list = []
    signals = []
    day = int(timestamp[index].split("/")[0])
    month = int(timestamp[index].split("/")[1])
    year = int(timestamp[index].split("/")[2])
    if month < 7:
        quarter = "Q2"
    else:
        quarter = "Q4"
    current = str(year) + quarter

    # Q-2 quarter
    if quarter == "Q2":
        year -= 1
        quarter = "Q4"
    else:
        quarter = "Q2"
    date = str(year) + quarter

    # Check if quarter is different from previous day's quarter
    quarter_list.append(current)
    if index == 0:
        store_index.append([index,current])
    else:
        if quarter_list[index] != quarter_list[index-1]:
            store_index.append([index,current])

    # Search for index for ranking
    for index2 in range(len(date_ranking)):
        if date == date_ranking[index2]:
            answer = index2

    # Search for index for market cap
    for index3 in range(len(date_marketcap)):
        if date == date_marketcap[index3]:
            answer_marketcap = index3

    # Loop through stocks
    for stock in range(len(ticker)):
        signal = 0
        if sector[stock] == "Consumer Discretionary":
            signal = int(ranking[answer][1])
            signals.append(signal)
        if sector[stock] == "Consumer Staples":
            signal = int(ranking[answer][2])
            signals.append(signal)
        if sector[stock] == "Energy":
            signal = int(ranking[answer][3])
            signals.append(signal)
        if sector[stock] == "Financials":
            signal = int(ranking[answer][4])
            signals.append(signal)
        if sector[stock] == "Health Care":
            signal = int(ranking[answer][5])
            signals.append(signal)
        if sector[stock] == "Industrials":
            signal = int(ranking[answer][6])
            signals.append(signal)
        if sector[stock] == "Information Technology":
            signal = int(ranking[answer][7])
            signals.append(signal)
        if sector[stock] == "Materials":
            signal = int(ranking[answer][8])
            signals.append(signal)
        if sector[stock] == "Real Estate":
            signal = int(ranking[answer][9])
            signals.append(signal)
        if sector[stock] == "Telecommunication Services":
            signal = int(ranking[answer][10])
            signals.append(signal)
        if sector[stock] == "Utilities":
            signal = int(ranking[answer][11])
            signals.append(signal)

        # Fill the result matrix with the return if signal = 1, otherwise zero
        if index == 0:  # Skip the first row
            resultMatrix[index + 1][stock + 1] = 0
        else:
            gain = functions.percentageChange(float(price[index-1][stock+1]),float(price[index][stock+1]))
            resultMatrix[index + 1][stock + 1] = gain*signal

        # Calculate number of stocks to be invested
        if signal == 1 and float(price[index][stock+1]) != 0:
            stockcount += 1
            market_cap_value = float(market_cap[answer_marketcap][stock + 1])
            market_cap_list.append(market_cap_value)
            market_cap_total += market_cap_value
        else:
            market_cap_list.append(0)

    # Calculate Daily Return for each loop
    row_to_calculate = resultMatrix[index+1][1:len(ticker)+1]  # take the row
    if stockcount == 0:
        resultMatrix[index + 1][len(ticker) + 1] = 0
    else:
        if market_cap_total == 0:
            return_today = sum(row_to_calculate)/stockcount
        else:
            for column in range(len(row_to_calculate)):
                weight = market_cap_list[column]/market_cap_total
                row_to_calculate[column] *= weight
            resultMatrix[index + 1][1:len(ticker) + 1] = row_to_calculate
            return_today = sum(row_to_calculate)
        resultMatrix[index + 1][len(ticker) + 1] = return_today
        daily_return.append(return_today)

    # Calculate Cumulative Return
    if index > 1:  # Skip First Row because all the prices are zero
        portfolio = float(resultMatrix[index][len(ticker) + 4]) * (1 + return_today/100)
        resultMatrix[index + 1][len(ticker) + 4] = portfolio
        portfolio_list.append(portfolio)
        cumulative_return = (portfolio - initial_amount) * 100 / initial_amount
        resultMatrix[index + 1][len(ticker) + 2] = cumulative_return
        cumulative_list.append(cumulative_return)
    numberofday += 1
    # Stock Count
    resultMatrix[index + 1][len(ticker) + 3] = stockcount

# Calculate Quarterly Return (HK: semiannually)
quarter_matrix = [[0 for x in range(len(store_index))] for y in range(2)]
quarter_matrix[0][0] = "Index"
quarter_matrix[-1][0] = "Quarterly return"
for i in range(len(store_index)-1):
    quarter_matrix[0][i+1] = store_index[i][1]
    quarter_matrix[1][i+1] = functions.percentageChange(portfolio_list[store_index[i][0]], portfolio_list[store_index[i+1][0]])
Quarter = quarter_matrix[0][1:]
Quarter_return = quarter_matrix[1][1:]

# Draw quarterly return bar graphs
fig= plt.figure()  # Create a figure
ax = plt.axes()  # Create axis
plt.bar(Quarter, Quarter_return, align='center', alpha=0.5)  # Draw Bar graph
ax.xaxis.set_major_locator(plt.MaxNLocator(8))  # Set Maximum number of x-axis values to show
fig.autofmt_xdate()  # Rotate values to see more clearly
plt.ylabel('%')
plt.title('Semiannual Return')
fig.savefig(directory + market + " Semiannual Return.png")
plt.show()

# Take the last index to get cumulative_return
cumulative_return = resultMatrix[-1][len(ticker)+2]  # Final cumulative return
# Calculate related stats
print(numberofday, "days;", round(cumulative_return,4),"%")
print("\n")
print("Average Daily Return =", round(average(daily_return), 4), "%")
print("SD of Daily Return =", round(SD(daily_return), 4), "%")
print("Annualized Return =", round(average(daily_return)*oneyear, 4), "%")
print("Annualized Volatility =", round(SD(daily_return)*(oneyear**(1/2)), 4), "%")
print("Annual Sharpe Ratio: ", round(sharpe(daily_return), 4))
print("Max Consecutive Loss = ", maxconsecutive(daily_return), "days")
print("Win Rate:", round(winrate(daily_return), 4), "%")
print("Max Drawdown:", round(functions.findMin(daily_return), 4), "%")
print("\n")

# Calculate stats for market index
daily_return_index = functions.percentagelist(price_index)
cumulative_index = functions.percentageChange(price_index[0], price_index[-1])
print(market,"Index: ")
print("Cumulative =", round(cumulative_index, 4), "%")
print("Average Daily Return =", round(average(daily_return_index), 4), "%")
print("SD of Daily Return =", round(SD(daily_return_index), 4), "%")
print("Annualized Return =", round(average(daily_return_index)*oneyear, 4), "%")
print("Annualized Volatility =", round(SD(daily_return_index)*(oneyear**(1/2)), 4), "%")
print("Annual Sharpe Ratio: ", round(sharpe(daily_return_index),4))
print("Max Consecutive Loss = ", maxconsecutive(daily_return_index), "days")
print("Win Rate:", round(winrate(daily_return_index), 4), "%")
print("Max Drawdown:", round(functions.findMin(daily_return_index), 4), "%")

# Find correct date to match dates for market index and our portfolio
correctdate = []
correctdate2 = []
for i in range(len(timestamp)):
    for j in range(len(date_index)):
        if date_index[j] == timestamp[i]:
            correctdate.append(i)
            correctdate2.append(j)

# Prepare for a graph
x_axis = date_index
y = []  # Strategy Portfolio
y_index = []  # Market Index Portfolio

# Normalize to initial amount and fill up values for y-axis by searching for correct dates
relative_portfolio = functions.relative(portfolio_list)
for i in range(len(correctdate)):
    correctindex = correctdate[i]-1
    y.append(relative_portfolio[correctindex])
for j in range(len(correctdate2)):
    correctindex = correctdate2[j]
    y_index.append(price_index[correctindex])
y_index = functions.relative(y_index)

# Plot a graph
fig = plt.figure()  # Create a figure
ax = plt.axes()  # Create axis
if len(x_axis) == len(y):
    plt.plot(x_axis, y, label='Strategy')
    plt.plot(x_axis, y_index, label=market)
    ax.xaxis.set_major_locator(plt.MaxNLocator(5)) # Set Maximum number of x-axis values to show
    fig.autofmt_xdate() # Rotate values to see more clearly
    legend = plt.legend(loc='best')
    title = "Sector Selection: Growth of "+str(initial_amount/1000000)+" million"
    plt.title(title)
    plt.ylabel("Cumulative Return")
    fig.savefig(directory + market + " Sector Selection.png")
    plt.show()

# Write csv file of result matrix inside result folder
with open(directory + 'result matrix' + '.csv', "w") as output:
    writer = csv.writer(output, lineterminator='\n')
    for val in resultMatrix:
        writer.writerow(val)
