from labels import Label
from constants import CHUNK_SIZE, MAX_ROWS
from typing import Iterator, cast
import pyarrow.parquet as pq
import pandas as pd

def default_labelling(article_type:str) -> Label:
    if article_type in ["reliable", "political", "clickbait"]:
        return Label.REAL
    return Label.FAKE

def only_fake_labelling(article_type:str) -> Label:
    if article_type == "fake":
        return Label.FAKE
    return Label.REAL

def not_reliable_labelling(article_type:str) -> Label:
    if article_type == "reliable":
        return Label.REAL
    return Label.FAKE

def LIAR_labelling(article_type:str) -> Label:
    if article_type in ["true", "half-true", "barely-true", "mostly-true"]:
        return Label.REAL
    return Label.FAKE

# Deprecated, don't use, just use pd.read_parquet instead
def dataset_iterator(dataset_file:str, columns:list[str] | None = None) -> Iterator[pd.DataFrame]:
    pq_file = pq.ParquetFile(dataset_file)
    rows_read = 0
    for batch in pq_file.iter_batches(batch_size=CHUNK_SIZE, columns=columns): # type: ignore
        rows_read += len(batch) # type: ignore
        if rows_read > MAX_ROWS and MAX_ROWS > 0:
            return
        # cast to ignore type warnings.
        yield cast(pd.DataFrame, batch.to_pandas()) # type: ignore

def csv_to_parquet(input_path: str, output_path: str):
    # csv.field_size_limit(sys.maxsize) # if i don't use this and engine="python" i cant read all datasets.
    # writer = None
    # for chunk in pd.read_csv(input_path, chunksize=CHUNK_SIZE, dtype=str, engine="python"): # who the fuck put a string in the id column
    #     chunk = chunk.fillna("")
    #     table = pa.Table.from_pandas(chunk) # type: ignore
    #     if writer is None:
    #         writer = pq.ParquetWriter(output_path, table.schema) # type: ignore
    #     writer.write_table(table) # type: ignore
    # writer.close() # type: ignore
    pd.read_csv(input_path, low_memory=False).to_parquet(output_path)

def get_time_boundaries (filename: str) : #type: ignore 
    # Only load the timestamp column to save RAM
    df_dates = pq.read_table(filename, columns=['scraped_at']).to_pandas() #type: ignore 
    df_dates['scraped_at'] = pd.to_datetime(df_dates['scraped_at'], format='ISO8601', errors='coerce',utc=True) #type: ignore 
    # Sort the dates to find the percentiles
    sorted_dates = df_dates['scraped_at'].sort_values() #type: ignore 
    # Find the timestamps at the 80th and 90th percentile
    train_cut = sorted_dates.quantile(0.80) #type: ignore 
    val_cut = sorted_dates.quantile(0.90) #type: ignore 
    return train_cut, val_cut #type: ignore 

