import pandas as pd
import random


def get_batches(indices, batch_size):
    for i in range(0, len(indices), batch_size):
        yield indices[i:i + batch_size]


# This one is used for prediction of the ongoing results of elections
def simulate_vote_result_batching(df: pd.DataFrame, number_of_batches: int) -> list[pd.DataFrame]:
    """ Simulate the process of batching vote results"""
    if number_of_batches > len(df):
        raise ValueError("Number of batches cannot be larger than the number of rows in the dataframe")

    batch_size = len(df) // number_of_batches
    index = df.index.array
    dataframes = [df.copy()]
    indices_to_batch = index.copy()
    # random.seed(420)
    random.shuffle(indices_to_batch)
    batched_indices = []

    for batch in get_batches(indices_to_batch, batch_size):
        copied_df = df.copy()
        batched_indices.extend(batch)
        copied_df.loc[batched_indices] = -1
        dataframes.append(copied_df)
    return list(reversed(dataframes))


def simulate_partial_vote_result(df: pd.DataFrame, min_v: float, max_v: float) -> pd.DataFrame:
    """
    Simulate a partial vote result
    Args:
        df (pd.DataFrame): The dataframe to simulate the result for

    Returns:
        pd.DataFrame: The simulated dataframe
        :param min_v: The minimum value to multiply the data by
        :param max_v: The maximum value to multiply the data by
    """
    df_n = df.copy()
    for i in df.index:
        dec = random.uniform(min_v, max_v)
        df_n.loc[i] = (df_n.loc[i] * dec) // 1
    return df_n


def get_samples_for_batching(dataframe: pd.DataFrame) -> list[tuple[pd.DataFrame, pd.DataFrame]]:
    """
    Get the samples for batching
    Returns:
        list[tuple[pd.DataFrame, pd.DataFrame]]: A list of tuples containing two dataframes. The first dataframe contains
        'historical data', while the second contains the data which will simulate the ongoing election.

    Attributes:
        dataframe (pd.DataFrame): The dataframe to get samples from.
        It is assumed that the dataframe contains all voting results on which testing is done.
    """
    samples = []

    for i in range(1, len(dataframe.columns)):
        samples.append((dataframe.iloc[:, :i], dataframe.iloc[:, i:i + 1]))
    return samples


def create_sample_dataframe(num_rows: int):
    # Create random data
    data = {
        'VoterID': range(0, num_rows),
        'Vote': [random.choice([1, 2]) for _ in range(num_rows)]
    }
    df = pd.DataFrame(data)
    df.set_index('VoterID', inplace=True)

    return df


def merge_historical_data(dfs: list[pd.DataFrame]) -> pd.DataFrame:
    """
    Merge the historical dataframes into one dataframe
    Args:
        dfs (list[pd.DataFrame]): The dataframes to merge

    Returns:
        pd.DataFrame: The merged dataframe
    """
    return pd.concat(dfs, axis=1)


def accumulate_results(results, div=False):
    total_votes = [0, 0, 0, 0]
    total_count = 0

    for result in results:
        if isinstance(result.count, pd.Series):
            count = result.count.iloc[0]
        else:
            count = result.count

        total_count += count
        for i in range(4):
            total_votes[i] += result.percentages[i] * count / 100

    if total_count == 0:
        return [0, 0, 0, 0]

    return [vote / total_count * 100 for vote in total_votes]