from datetime import datetime

import numpy as np
import polars as pl
from pathlib import Path
from polars import col, lit
import pandas as pd
from pandas.io.common import get_handle

pl.__version__

'0.7.18'

pd.__version__

'1.2.4'

# Download a huge csv as a test. Takes a while and only needed once...
big_csv = Path("./big.csv")
csv_url = "http://sdm.lbl.gov/fastbit/data/star2002-full.csv.gz"

if not big_csv.exists():
    with get_handle(csv_url, compression="gzip", mode="r") as fh_in, open(big_csv, "bw") as fh_out:
        fh_out.write(fh_in.handle.buffer.read())

edf = pl.read_csv(str(big_csv), has_headers=False)

edf.filter(col("column_1") == 1).select(["column_9"]).head()

edf[edf["column_1"] == 1][["column_9"]].head()

ldf = pl.scan_csv(str(big_csv), has_headers=False)

ldf = ldf.filter(col("column_1") == 1)
ldf.select(["column_9"]).collect().head()

ldf = pl.scan_csv(str(big_csv), has_headers=False)
ldf[ldf["column_1"] == 1][["column_9"]].head()

---------------------------------------------------------------------------

TypeError                                 Traceback (most recent call last)

<ipython-input-10-aaa4e676d17c> in <module>
      1 ldf = pl.scan_csv(str(big_csv), has_headers=False)
----> 2 ldf[ldf["column_1"] == 1][["column_9"]].head()


TypeError: 'LazyFrame' object is not subscriptable

edf[1] # this is a bug right now

Series: 'column_2' [i64]
[
    1613423
    1613423
    1613423
    1613423
    1613423
    1613423
    1613423
    1613423
    1613423
    1613423
]

edf[[1]]

edf[1, 3] # index by (row, column)

20011015.222604

edf[1, "column_4"] # or as string

20011015.222604

edf[1, [2, 3]]  # index by (row, column) but returns data frame

edf[1:4, "column_4":] # slice by row and column's name

pdf = edf.to_pandas()

pdf.iloc[1, 3]

20011015.222604

# when mixing indexing by integer and by string it gets less comprehensible in Pandas
pdf["column_4"].iloc[1] 

20011015.222604

pdf.iloc[1, [2, 3]]

column_3    8.080000e+02
column_4    2.001102e+07
Name: 1, dtype: float64

pdf[["column_4"]].iloc[1:4]

# for slicing with column names and guaranteed indexing by integer we have to write:
pdf.loc[:, "column_4":].iloc[1:4]
# `pdf.loc[1:4, "column_4":]` works only as long the index is set correctly.

left_df = pl.DataFrame({"a": [1, 2, 3], "b": [None, "b", "c"]})
right_df = pl.DataFrame({"a": [1, 2], "c": [42, 69]})

df = left_df.join(right_df, on="a", how="left")
df

df.filter(col("c").is_null())

left_pdf = left_df.to_pandas()
right_pdf = right_df.to_pandas()

right_pdf.dtypes

a    int64
c    int64
dtype: object

pdf = pd.merge(left_pdf, right_pdf, on="a", how="left")

pdf

df.with_column((lit(3)*col("c")).alias("3*c"))

pdf.assign(**{"3*c": 3*pdf["c"]})

df = pl.DataFrame(
    {
        "nrs": [1, 2, 3, None, 5],
        "names": ["foo", "ham", "spam", "egg", None],
        "random": np.random.rand(5),
        "groups": ["A", "A", "B", "C", "B"],
    }
)
df

# and in Pandas
pdf = df.to_pandas()

df.select(
    [
        pl.sum("nrs"), # or equivalently col("nrs").sum()
        col("names").sort(),
        col("names").n_unique().alias("unique_names_1"),
    ]
)

pd.DataFrame(
    {
        "nrs": pdf["nrs"].sum(),
        "names": pdf["names"].sort_values(),
        "unique_names_1": pdf["names"].nunique(dropna=False),
    }
)

df.select(col("names").filter(col("random") > 0.4))

pdf.loc[pdf["random"] > 0.4][["names"]]

pdf.query("random > 0.4")[["names"]]

df.select(
    [
        (pl.sum("nrs") * pl.when(col("random") > 0.5)
                           .then(0)
                           .otherwise(col("random"))
        ).alias("result")
    ]
)

pd.Series(
    np.where(pdf["random"] > 0.5, 0, pdf["random"] * pdf["nrs"].sum()), name="result"
).to_frame()

# or easier to read but much slower since not vectorized at all
pdf.random.apply(lambda x: 0 if x > 0.5 else x * pdf["nrs"].sum()).to_frame()

df.select(
    [
        col("*"),  # select all
        col("random").sum().over("groups").alias("sum[random]/groups"),
        col("random").list().over("names").alias("random/name"),
    ]
)

(pdf.set_index("groups")
    .assign(**{"sum[random]/groups": pdf.groupby(['groups'])['random'].sum()})
    .set_index("names")
    .assign(**{"random/name": pdf.groupby(['names'], dropna=False)['random'].apply(list)})
    .reset_index()
)

# or alternatively using `join` which is also avoiding the NaN problem
pdf.join(
    pdf.groupby('groups').random.sum().rename("sum[random]/groups"), on="groups"
).join(
    pdf.groupby('names', dropna=False).random.apply(list).rename("random/name"), on="names"
)

df = pl.read_csv("https://theunitedstates.io/congress-legislators/legislators-current.csv")
pdf = df.to_pandas()

(df.lazy()  # allows for working only on a subset using limit
   .groupby("first_name")
   .agg(
       [
           col("party").count().alias("n_party"),  # renaming an aggregated column is a bliss
           col("gender").list(),
           col("last_name").first(),
       ]
   )
   .sort("n_party", reverse=True)
   .limit(5)
   .collect()
)

(pdf.groupby("first_name")
    .agg({"party": "count", 
          "gender": lambda grp: grp.to_list(), 
          "last_name": "first"})
    .rename(columns={"party": "n_party", 
                     "gender": "gender_agg_list", 
                     "last_name": "last_name_first"})
    .sort_values(by="n_party", ascending=False)
    .reset_index()
    .head(5))

(df.lazy()
   .groupby("state")
   .agg(
       [
           (col("party") == "Democrat").sum().alias("demo"),
           (col("party") == "Republican").sum().alias("rep"),
       ]
   )
   .sort("demo", reverse=True)
   .limit(5)
   .collect()
)

(pdf.groupby("state")
    .agg({"party": [("demo", lambda grp: np.sum(grp == "Democrat")), 
                    ("rep", lambda grp:  np.sum(grp == "Republican"))]})
    .droplevel(0, axis=1) 
    .reset_index()
    .sort_values(by="demo", ascending=False)
    .head(5)
)

def compute_age() -> pl.Expr:
    # Date64 is time in ms
    ms_to_year = 1e3 * 3600 * 24 * 365
    return (
        lit(datetime(2021, 1, 1)) - col("birthday")
    ) / (ms_to_year)


def avg_age(gender: str) -> pl.Expr:
    return (
        compute_age()
        .filter(col("gender") == gender)
        .mean()
        .alias(f"avg {gender} age")
    )


(df.lazy()
   .groupby(["state"])
   .agg(
       [
            avg_age("M"),
            avg_age("F"),
       ]
   )
   .sort("state")
   .limit(5)
   .collect()
)

def p_compute_age(grp: pd.DataFrame):
    # Date64 is time in ms
    s_to_year = 3600 * 24 * 365
    return (
        (datetime(2021, 1, 1) - grp["birthday"]).dt.total_seconds()
    ) / (s_to_year)


def p_avg_age(grp: pd.DataFrame, gender: str):
    age = p_compute_age(grp)
    mean_age = age[grp["gender"] == gender].mean()
    return pd.Series([mean_age], index=[f"avg {gender} age"])


(pdf.groupby("state")
    .apply(
        lambda grp: pd.concat(
            [p_avg_age(grp, gender="M"),
             p_avg_age(grp, gender="F")]
        )
    )
    .reset_index()
    .sort_values(by="state")
    .head(5)
)

df = pl.DataFrame({"foo": np.arange(10), 
                   "bar": np.random.rand(10), 
                   "cls": np.random.randint(2, size=10)})
pdf = df.to_pandas()

df

def my_custom_func(s: pl.Series) -> pl.Series:
    return np.exp(s) / np.log(s)

df.filter(col("bar").map(my_custom_func, return_dtype=pl.Float64) > -1)

df.select(col("bar").apply(lambda x: 3*x))

pdf[pdf["bar"].apply(my_custom_func) > -1]

pdf[["bar"]].applymap(lambda x: 3*x)

df.groupby(["cls"]).agg([col("bar").apply(lambda grp: 3*grp.sum())])

pdf.groupby("cls", as_index=False).agg({"bar": lambda grp: 3 * grp.sum()})