from input import Input
from word import Word
from data import Data
from display import Display as dp
from colorama import Fore
import numpy as np

# The Game Object
class LeWord:

    # Only gets the mode,
    def __init__(self, mode):
        self.mode = mode

    # The joker of a word is the count of vowels and consonants in it.
    def joker(self, true_word, mode=1):
        vowels = "AEIOUaeiou"

        # checks the mode
        if mode == 1:
            vowel_c = 0
            consonant_c = 0

            # for every letter if it is vowel vowel count plus one else consonants plus one
            for letter in true_word:
                if letter in vowels:
                    vowel_c += 1
                    continue
                consonant_c += 1

            # then displays the joker for the word
            dp.display_joker(vowel_c, consonant_c, mode=mode)
        else:
            # if mode is not 1 then for every word does the same process
            for index, word in enumerate(true_word):
                vowel_c = 0
                consonant_c = 0

                # for every letter if it is vowel vowel count plus one else consonants plus one
                for letter in word:
                    if letter in vowels:
                        vowel_c += 1
                        continue
                    consonant_c += 1

                # then displays the joker for the word
                dp.display_joker(vowel_c, consonant_c, mode=mode, que=index+1)
        print()

    # starts the game with selected mode
    def play_game(self, again):

        if self.mode[2] in ["1", "one"]:
            self.single_word_game()
        else:
            self.multi_word_game()

    # the game with one-single word
    def single_word_game(self):

        # Picks the random word
        true_word = Word.pick_random_word(self.mode[1])

        # Turn true word into array of individual letters and indexes using numpy library
        true_arr = np.array(list(true_word.upper()))

        # Turn the array into [item id found] type
        true_idx = [[item, idx, None] for idx, item in enumerate(true_arr)]

        attempt = 0
        guesses = []

        # while the attempt count is less than the word length + 1
        while attempt < self.mode[1] + 1:
            # gets the guess from the user
            guess_word = Input.ask_user_input(self.mode[1], attempt)

            # looks for special inputs
            if not guess_word:
                break
            elif guess_word == "JOKER":
                self.joker(true_word, mode=self.mode[3])
                continue
            else:
                # then appends the guess
                guesses.append(guess_word)

                # Turn guess into array of individual letters and indexes using numpy library
                guess_arr = np.array(list(guess_word))

                # Turn the array into [item id found] type
                guess_idx = [[item, idx, None] for idx, item in enumerate(guess_arr)]

                matched = []
                existing = []
                # uses numpy to look if any matches between true and guess array pairs then appends to 'matching'
                matching = np.where(true_arr == guess_arr)[0]

                # for the matches between two array puts 'YES' to found
                for item in matching:
                    matched.append(guess_idx[item][0])
                    guess_idx[item][2], true_idx[item][2] = 'YES', 'YES'

                # gets the unfounded words
                rem_guess = [item for item in guess_idx if item[2] != 'YES']
                rem_true = [item for item in true_idx if item[2] != 'YES']

                # for the letters that unfounds looks if they are exist in somewhere in true word
                # if it is then marks found with 'EX'
                for guess in rem_guess:
                    for true in rem_true:
                        if guess[0] == true[0]:
                            if list(true_arr).count(guess[0]) > (matched.count(guess[0]) + existing.count(guess[0])):
                                existing.append(guess[0])
                                guess[2], true[2] = 'EX', 'EX'
                            else:
                                continue

                # Colors and marks letters based on match, exist, or not exist
                final = Word.mark_letters(guess_idx)

                # adds space to make it identical in table
                for letter in final:
                    letter += " "

                # Turn current guess into table
                guess_df = Word.build_df(attempt, final, attempt)

                print(guess_df)

                # checks whether all the letters are found or not.
                if guess_word.lower() == true_word.lower():
                    # if it is then makes the congrats
                    dp.display_cong(1, true_word=true_word.upper())
                    # and asks the user want to play again
                    again = True if input(Fore.RED + "-->" + Fore.RESET).lower() == "y" else False
                    imp = Input.check_input(str(again))
                    if imp == "m":
                        return False
                    print()
                    # finally writes the data into the csv file
                    Data.write_game_data([self.mode[2], true_word, len(true_word), attempt + 1, guesses,
                                          True if len(true_word) > attempt + 1 else False])

                    return again
                else:
                    # if not plus ones the attempt count and continues
                    attempt += 1

                    # if attempt count is at limit then
                    if attempt == self.mode[1] + 1:
                        # gives the fail message
                        dp.display_fail(1, true_word)
                        # and asks the user want to play again
                        again = True if input(Fore.RED + "-->" + Fore.RESET).lower() == "y" else False
                        imp = Input.check_input(str(again))
                        if imp == "m":
                            return False

                        print()
                        # finally writes the data into the csv file
                        Data.write_game_data([self.mode[2], true_word, len(true_word), attempt + 1, guesses,
                                              True if len(true_word) > attempt + 1 else False])
                        return again

    def multi_word_game(self):
        founded_words = []

        # Picks the random words
        true_word = Word.pick_random_word(self.mode[1], int(self.mode[3]))

        # Turn true words into array of individual letters and indexes using numpy library
        true_arr = np.array(list([j for j in i.lower()] for i in true_word))

        # Turn the array into [item id found] type
        true_idx = [[[item, idx, None] for idx, item in enumerate(element)] for element in true_arr]

        attempt = 0
        guesses_word = []

        # while the attempt count is less than the word length + 5
        while attempt < self.mode[1] + 5:
            # gets the guess from the user
            guess_word = Input.ask_user_input(self.mode[1], attempt)

            # looks for special inputs
            if not guess_word:
                break
            elif guess_word == "JOKER":
                self.joker(true_word, mode=self.mode[3])
                continue
            else:
                # then appends the guess
                guesses_word.append(guess_word)

                # Turn guess into array of individual letters and indexes using numpy library
                guess_arr = np.array(list(guess_word.lower()))

                # Turn the array into [item id found] type
                guess_idx = [[[item, idx, None] for idx, item in enumerate(guess_arr)] for i in range(int(self.mode[3]))]

                matched = []
                existing = []
                matches = []

                # for every word in true words
                for element in true_arr:
                    # uses numpy to look if any matches between true and guess array pairs then appends to 'matched'
                    matches.append(np.where(element == guess_arr)[0])

                # looks if other words founded before if it is then displays them
                for founds in founded_words:
                    index, word = founds[0], founds[1]
                    guess_idx[index] = [[item, idx, 'YES'] for idx, item in enumerate(word)]

                # for the matches between two array puts 'YES' to found
                for index, match in enumerate(matches):
                    for item in match:
                        matched.append([index, guess_idx[index][item][0]])
                        guess_idx[index][item][2], true_idx[index][item][2] = 'YES', 'YES'

                # gets the unfounded letters
                rem_guess = []
                for element in guess_idx:
                    word = []
                    for item in element:
                        if item[2] != 'YES':
                            word.append(item)
                    rem_guess.append(word)

                # gets the unfounded letters
                rem_true = []
                for element in true_idx:
                    word = []
                    for item in element:
                        if item[2] != 'YES':
                            word.append(item)
                    rem_true.append(word)

                # for the letters that unfounds looks if they are exist in somewhere in true word
                # if it is then marks found with 'EX'
                for index_g, guesses in enumerate(rem_guess):
                    for guess in guesses:
                        for index, word in enumerate(rem_true):
                            for true in word:
                                if index_g == index:
                                    if guess[0] == true[0]:
                                        if list(true_arr[index]).count(guess[0]) > \
                                                (matched.count(guess[0]) + existing.count(guess[0])):
                                            existing.append(guess[0])
                                            guess[2], true[2] = 'EX', 'EX'
                                        else:
                                            continue

                # Colors and marks letters based on match, exist, or not exist
                final = Word.mark_letters(guess_idx, mode=self.mode[3])

                # adds space to make it identical in table
                for letter in final:
                    letter += " "

                # Turn current guess into table
                guess_df = Word.build_df(attempt, final, attempt, mode=int(self.mode[3]))

                print(guess_df)

                # if any word is founded then adds them to founded words
                for index, word in enumerate(true_word):
                    if guess_word.lower() == word.lower():
                        founded_words.append([index, word])

                # checks whether all the letters are found or not.
                if len(true_word) == len(founded_words):
                    # if it is then makes the congrats
                    dp.display_cong(len(true_word), true_word=true_word)
                    # and asks the user want to play again
                    again = True if input(Fore.RED + "-->" + Fore.RESET).lower() == "y" else False
                    imp = Input.check_input(str(again))
                    if imp == "m":
                        return False

                    print()
                    # finally writes the data into the csv file
                    Data.write_game_data([self.mode[2], true_word, len(true_word), attempt + 1, guesses_word,
                                          True if len(true_word) > attempt + 1 else False])
                    return again
                else:
                    # if not plus ones the attempt count and continues
                    attempt += 1

                    # if attempt count is at limit then
                    if attempt == self.mode[1] + 5:
                        # gives the fail message
                        dp.display_fail(len(true_word), true_word)
                        # and asks the user want to play again
                        again = True if input(Fore.RED + "-->" + Fore.RESET).lower() == "y" else False
                        imp = Input.check_input(str(again))
                        if imp == "m":
                            return False

                        print()
                        # finally writes the data into the csv file
                        Data.write_game_data([self.mode[2], true_word, len(true_word), attempt + 1, guesses_word,
                                              True if len(true_word) > attempt + 1 else False])
                        return again

    # Starts the LE-WORD
    @staticmethod
    def start_le_word():

        opening = True
        again = False
        # Displays the menu for user to choose
        mode = dp.display_menu()

        # Until the user types mode 'quit' the game continues
        while mode[0] != "quit":
            # if it is not the first time and not the again after the game prints menu
            if not opening and not again:
                mode = dp.display_menu()
                if mode[0] == "quit":
                    break
            # if choice statics displays statics
            if mode[0] == "statics":
                dp.display_statics()
            # if choice rules displays rules
            elif mode[0] == "rules":
                dp.display_rules()

            opening = False

            # Creates a LeWord object
            game = LeWord(mode)

            # if the mode is play
            if mode[0] == "play":
                # then looks at the game mode
                if mode[2] in ["one", "1"]:
                    # if single then single_word_game
                    again = game.single_word_game()
                elif mode[2] == "multi":
                    # if multi then multi_word_game
                    again = game.multi_word_game()

                # after game ends if user wants to play again
                if again:
                    # asks user to play in which mode
                    a, b, c = Input.ask_mode()
                    if not mode:
                        opening = False
                        again = False
                    # assigns the mode and continues
                    mode = ["play", a, b, c]


# __author__     = "Chang Gao"
# __copyright__  = "Copyright 2018 to the author"
# __license__    = "Private"
# __version__    = "0.1.0"
# __maintainer__ = "Chang Gao"
# __email__      = "chang.gao@uzh.ch"
# __status__     = "Prototype"
import sys
import os
import torch as t
import torch.nn.functional as F
from torch.autograd.function import Function
import time
import math

def save_normalization(save_path, tr_mean, tr_std, lab_mean, lab_std):
    fn_base = os.path.splitext(save_path)[0]
    print("\nSaving normalization parameters to " + str(fn_base)+'-XX.pt')
    norm = {
            'tr_mean': tr_mean,
            'tr_std': tr_std,
            'lab_mean': lab_mean,
            'lab_std': lab_std,
        }
    t.save(norm, str(fn_base+'-norm.pt'))

def load_normalization(save_path):
    fn_base = os.path.splitext(save_path)[0]
    print("\nLoading normalization parameters from ", str(fn_base))
    norm = t.load(fn_base+'-norm.pt')
    return norm['tr_mean'], norm['tr_std'], norm['lab_mean'], norm['lab_std']

# print command line (maybe to use in a script)
def print_commandline(parser):
    args = parser.parse_args()
    print('Command line:')
    print('python '+os.path.basename(sys.argv[0]), end=' ')
    for arg in vars(args):
        print('--' + str(arg) + ' "' + str(getattr(args, arg))+'"', end=' ')
    print()

def timeSince(since):
    now = time.time()
    s = now - since
    m = math.floor(s / 60)
    s -= m * 60
    return '%dm %ds' % (m, s)

def to_sparse(x):
    """ converts dense tensor x to sparse format """
    x_typename = t.typename(x).split('.')[-1]
    sparse_tensortype = getattr(t.sparse, x_typename)

    indices = t.nonzero(x)
    if indices.nelement() == 0:  # if all elements are zeros
        print("1", indices)
        return sparse_tensortype(*x.shape)
    else:
        print("2", indices)
        indices = indices.t()
        values = x[tuple(indices[i] for i in range(indices.shape[0]))]
        return sparse_tensortype(indices, values, x.size())

def quantizeTensor(x, m, n, en):
    """
    :param x: input tensor
    :param m: number of integer bits before the decimal point
    :param n: number of fraction bits after the decimal point
    :return: tensor quantized to fixed-point precision
    """
    if en == 0:
        return x
    power = 2. ** n
    clip_val = 2. ** (m + n - 1)
    value = t.round(x * power)
    # value = GradPreserveRoundOp.apply(x * power)  # rounding
    value = t.clamp(value, -clip_val, clip_val - 1)  # saturation arithmetic
    value = value / power
    return value

def quantize_rnn(net, qi, qf, en):
    for name, param in net.named_parameters():
        if 'rnn' in name:
            param.data = quantizeTensor(param.data, qi, qf, en)
    return net

def pruneTensor(x, alpha):
    """
    :param x: input tensor
    :param m: number of integer bits before the decimal point
    :param n: number of fraction bits after the decimal point
    :return: tensor quantized to fixed-point precision
    """

    n_neuron = x.size(0)
    n_input = x.size(1)
    prune_prob_mask = t.exp(-alpha * t.unsqueeze(t.arange(0, n_neuron), dim=1).repeat(1, n_input).float()).cuda()
    prune_rand_mask = t.rand(n_neuron, n_input).cuda()
    prune_mask = prune_rand_mask.masked_fill_(prune_rand_mask > prune_prob_mask, 1)
    prune_mask = prune_mask.masked_fill_(prune_rand_mask <= prune_prob_mask, 0)
    _, indices = t.sort(t.abs(x), 0)
    # print("indices shape", indices.size())
    # print("prune_mask shape", prune_mask.size())
    # print("x shape", x.size())
    for j in range(0, n_input):
        x[indices[:, j], j] *= prune_mask[:, j]

    return x

def targetedDropout(x, gamma, alpha, epoch):
    """
    :param x: input tensor
    :param m: number of integer bits before the decimal point
    :param n: number of fraction bits after the decimal point
    :return: tensor quantized to fixed-point precision
    """

    t.manual_seed(epoch)
    t.cuda.manual_seed_all(epoch)

    n_elements = x.numel()
    drop_part = round(n_elements * gamma)
    weight_vec = x.view(-1)
    weight_vec_abs = t.abs(weight_vec)
    sorted, indices = t.sort(weight_vec_abs)
    # print(sorted)
    drop_indices = indices[0:drop_part]
    drop_rand_mask = t.rand(drop_indices.size(0)).cuda()
    drop_mask = t.ones(drop_indices.size(0)).cuda()
    drop_mask = drop_mask.masked_fill_(drop_rand_mask <= alpha, 0)
    weight_vec[drop_indices] *= drop_mask
    weight = t.reshape(weight_vec, (x.size(0), x.size(1)))

    return weight

def alignedTargetedDropout(x, gamma, alpha, num_pe, epoch):
    """
    :param x: input tensor
    :param m: number of integer bits before the decimal point
    :param n: number of fraction bits after the decimal point
    :return: tensor quantized to fixed-point precision
    """

    n_rows = x.shape[0]
    n_cols = x.shape[1]

    # Split and shuffle weight matrix
    for i in range(0, num_pe):
        for j in range(0, n_cols):
            targetedDropout(x[np.arange(i, n_rows, num_pe), j], gamma, alpha, epoch)
    return x

class GradPreserveRoundOp(Function):

    # Note that both forward and backward are @staticmethods
    @staticmethod
    def forward(ctx, input):
        output = t.round(input)
        return output

    # This function has only a single output, so it gets only one gradient
    @staticmethod
    def backward(ctx, grad_output):
        # This is a pattern that is very convenient - at the top of backward
        # unpack saved_tensors and initialize all gradients w.r.t. inputs to
        # None. Thanks to the fact that additional trailing Nones are
        # ignored, the return statement is simple even when the function has
        # optional inputs.
        grad_input = grad_output

        # These needs_input_grad checks are optional and there only to
        # improve efficiency. If you want to make your code simpler, you can
        # skip them. Returning gradients for inputs that don't require it is
        # not an error.
        # print(grad_output.size())
        # if not t.equal(grad_output, QuantizeT(grad_output, dW_qp)): print("grad_output not quantized")
        if ctx.needs_input_grad[0]:
            grad_input = grad_output

        # Return same number of parameters as "def forward(...)"
        return grad_input

class GradPreserveThreshold(Function):

    # Note that both forward and backward are @staticmethods
    @staticmethod
    def forward(ctx, input, threshold, value):
        output = F.threshold(input, threshold, value)
        return output

    # This function has only a single output, so it gets only one gradient
    @staticmethod
    def backward(ctx, grad_output):
        # This is a pattern that is very convenient - at the top of backward
        # unpack saved_tensors and initialize all gradients w.r.t. inputs to
        # None. Thanks to the fact that additional trailing Nones are
        # ignored, the return statement is simple even when the function has
        # optional inputs.
        grad_input = grad_output

        # These needs_input_grad checks are optional and there only to
        # improve efficiency. If you want to make your code simpler, you can
        # skip them. Returning gradients for inputs that don't require it is
        # not an error.
        # print(grad_output.size())
        # if not t.equal(grad_output, QuantizeT(grad_output, dW_qp)): print("grad_output not quantized")
        if ctx.needs_input_grad[0]:
            grad_input = grad_output

        # Return same number of parameters as "def forward(...)"
        return grad_input

def look_ahead_seq(seq_in, t_width=16, padding=0, batch_first=0):
    # Convert input sequence to batch first shape (seq_len, n_batch, n_feature)
    seq = seq_in
    if batch_first:
        seq = seq_in.transpose(0, 1)

    seq_len = seq.size(0)
    n_batch = seq.size(1)
    n_feature = seq.size(2)
    # int(t.ceil(float(seq_len)/float(t_width)))
    new_seq = []
    for i in range(0, seq_len):
        if i < seq_len - t_width:
            seq_block = seq[i:i + t_width, :, :]
        else:
            seq_block = seq[i:, :, :]
            seq_block_pad = t.zeros([t_width - (seq_len - i), n_batch, n_feature], dtype=t.float32).cuda()
            seq_block = t.cat((seq_block, seq_block_pad), 0)
        new_seq.append(seq_block)
    new_seq = t.stack(new_seq, 0)
    new_seq = new_seq.transpose(1, 2)
    new_seq = new_seq.transpose(0, 1)
    new_seq = new_seq.transpose(2, 3)
    return new_seq

def look_around_seq(seq_in, t_width=16, padding=0, batch_first=0):
    # Convert input sequence to batch first shape (seq_len, n_batch, n_feature)
    seq = seq_in
    if batch_first:
        seq = seq_in.transpose(0, 1)

    seq_len = seq.size(0)
    n_batch = seq.size(1)
    n_feature = seq.size(2)
    # int(t.ceil(float(seq_len)/float(t_width)))
    new_seq = []
    for i in range(0, seq_len):
        if i >= seq_len - t_width:
            seq_block = seq[i - t_width:, :, :]
            seq_block_pad = t.zeros([t_width - (seq_len - i) + 1, n_batch, n_feature], dtype=t.float32).cuda()
            seq_block = t.cat((seq_block, seq_block_pad), 0)
        elif i < t_width:
            seq_block = seq[0:i + 1 + t_width, :, :]
            seq_block_pad = t.zeros([t_width - i, n_batch, n_feature], dtype=t.float32).cuda()
            seq_block = t.cat((seq_block, seq_block_pad), 0)
        else:
            seq_block = seq[i - t_width:i + 1 + t_width, :, :]
        # print(seq_block.size())
        new_seq.append(seq_block)
    new_seq = t.stack(new_seq, 0)
    new_seq = new_seq.transpose(1, 2)
    new_seq = new_seq.transpose(0, 1)
    new_seq = new_seq.transpose(2, 3)
    return new_seq

def get_temporal_sparsity(list_layer, seq_len, threshold):
    # Evaluate Sparsity
    num_zeros = 0
    num_elems = 0
    # print(seq_len.size())
    # Iterate through layers
    for layer in list_layer:
        all_delta_vec = layer.transpose(0, 1)
        all_delta_vec = t.abs(all_delta_vec)  # Take absolute values of all delta vector elements
        for i, delta_vec in enumerate(all_delta_vec):
            seq = delta_vec[:seq_len[i], :]
            zero_mask = seq < threshold
            num_zeros += t.sum(zero_mask)
            num_elems += t.numel(zero_mask)
    sparsity = float(num_zeros) / float(num_elems)
    return sparsity


import pandas as pd
import yfinance as yf

# Read the symbols from a CSV file
symbols_df = pd.read_csv("symbols.csv")
symbols = symbols_df["Symbol"].tolist()

# Specify the years
years = [2021, 2022]

# Create an empty list to store the dataframes for each stock
dfs = []

# Iterate over the symbols
for symbol in symbols:
    try:
        # Get the dividends data for the specified years
        stock = yf.Ticker(symbol)
        dividends = stock.dividends

        for year in years:
            dividends_year = dividends.loc[str(year)]

            if dividends_year.empty:
                print(f"No dividend data available for {symbol} in {year}.")
                continue

            # Get the market capitalization data
            try:
                market_cap = float(stock.info["marketCap"])
            except KeyError:
                market_cap = None

            # Get the stock price data for the specified year
            stock_price = stock.history(start=f"{year}-01-01", end=f"{year}-12-31")

            # Calculate the price at the beginning and end of the year
            price_at_beginning = stock_price.iloc[0]['Close']
            price_at_end = stock_price.iloc[-1]['Close']

            # Create a dataframe for the current stock and year
            data = {
                "Symbol": symbol,
                "Year": year,
                "Dividend Date": dividends_year.index.strftime('%m/%d/%Y'),
                "Market Capitalization": market_cap,
                "Count of total dividends paid for that year": len(dividends_year),
                "How much was paid": dividends_year.tolist(),
                "Price at Beginning of Year": price_at_beginning,
                "Price at End of Year": price_at_end
            }
            df = pd.DataFrame(data)

            # Append the dataframe to the list
            dfs.append(df)

    except Exception as e:
        print(f"Got error from Yahoo API for ticker {symbol}, Error: {str(e)}")
        print(f"Skipping symbol {symbol} due to data unavailability.")

# Concatenate the dataframes for all stocks and years
result_df = pd.concat(dfs)

# Convert columns to float type
result_df["Market Capitalization"] = result_df["Market Capitalization"].astype(float)
result_df["Count of total dividends paid for that year"] = result_df["Count of total dividends paid for that year"].astype(float)

# Save the dataframe to a CSV file
result_df.to_csv("temp_dividend_data_2022_2023.csv", index=False)