import numpy as np
import pandas as pd
from cluster_experiments import PowerAnalysis

# Create fake data
N = 1_000
df = pd.DataFrame(
    {
        "target": np.random.normal(0, 1, size=N),
    }
)

config = {
    "analysis": "ols_non_clustered",
    "perturbator": "constant",
    "splitter": "non_clustered",
    "n_simulations": 50,
}
pw = PowerAnalysis.from_dict(config)

# Keep in mind that the average effect is the absolute effect added, this is not relative!
power = pw.power_analysis(df, average_effect=0.1)

# You may also get the power curve by running the power analysis with different average effects
power_line = pw.power_line(df, average_effects=[0, 0.1, 0.2])


# A faster method can be used to run the power analysis, using the approximation of
# the central limit theorem, which is stable with less simulations
from cluster_experiments import NormalPowerAnalysis
npw = NormalPowerAnalysis.from_dict(
    {
        "analysis": "ols_non_clustered",
        "splitter": "non_clustered",
        "n_simulations": 5,
    }
)
power_line_normal = npw.power_line(df, average_effects=[0, 0.1, 0.2])

# you can also use the normal power to get mde from a power level
mde = npw.mde(df, power=0.8)

from datetime import date

import numpy as np
import pandas as pd
from cluster_experiments.power_analysis import PowerAnalysis

# Create fake data
N = 1_000
clusters = [f"Cluster {i}" for i in range(100)]
dates = [f"{date(2022, 1, i):%Y-%m-%d}" for i in range(1, 32)]
df = pd.DataFrame(
    {
        "cluster": np.random.choice(clusters, size=N),
        "target": np.random.normal(0, 1, size=N),
        "date": np.random.choice(dates, size=N),
    }
)

config = {
    "cluster_cols": ["cluster", "date"],
    "analysis": "gee",
    "perturbator": "constant",
    "splitter": "clustered",
    "n_simulations": 50,
}
pw = PowerAnalysis.from_dict(config)

print(df)
# Keep in mind that the average effect is the absolute effect added, this is not relative!
power = pw.power_analysis(df, average_effect=0.1)
print(f"{power = }")

from datetime import date

import numpy as np
import pandas as pd
from cluster_experiments.experiment_analysis import GeeExperimentAnalysis
from cluster_experiments.perturbator import ConstantPerturbator
from cluster_experiments.power_analysis import PowerAnalysis, NormalPowerAnalysis
from cluster_experiments.random_splitter import ClusteredSplitter

# Create fake data
N = 1_000
clusters = [f"Cluster {i}" for i in range(100)]
dates = [f"{date(2022, 1, i):%Y-%m-%d}" for i in range(1, 32)]
df = pd.DataFrame(
    {
        "cluster": np.random.choice(clusters, size=N),
        "target": np.random.normal(0, 1, size=N),
        "date": np.random.choice(dates, size=N),
    }
)

# A switchback experiment is going to be run, prepare the switchback splitter for the analysis
sw = ClusteredSplitter(
    cluster_cols=["cluster", "date"],
)

# We use a constant perturbator to add artificial effect on the treated on the power analysis
perturbator = ConstantPerturbator()

# Use gee to run the analysis
analysis = GeeExperimentAnalysis(
    cluster_cols=["cluster", "date"],
)

# Run the power analysis
pw = PowerAnalysis(
    perturbator=perturbator, splitter=sw, analysis=analysis, n_simulations=50, seed=123
)

# Keep in mind that the average effect is the absolute effect added, this is not relative!
power = pw.power_analysis(df, average_effect=0.1)
print(f"{power = }")

# You can also use normal power analysis, that uses central limit theorem to estimate power, and it should be stable in less simulations
npw = NormalPowerAnalysis(
    splitter=sw, analysis=analysis, n_simulations=50, seed=123
)
power = npw.power_analysis(df, average_effect=0.1)
print(f"{power = }")

import numpy as np
import pandas as pd
from cluster_experiments import AnalysisPlan, SimpleMetric, Variant, Dimension

# ---------------------------------------
# ----------- Create fake data ----------
# ---------------------------------------

NUM_ORDERS = 10_000
NUM_CUSTOMERS = 3_000
EXPERIMENT_GROUPS = ["control", "treatment_1", "treatment_2"]
GROUP_SIZE = NUM_CUSTOMERS // len(EXPERIMENT_GROUPS)

# Generate customers and assign them to experiment groups
customer_ids = np.arange(1, NUM_CUSTOMERS + 1)
np.random.shuffle(customer_ids)
experiment_group = np.repeat(EXPERIMENT_GROUPS, GROUP_SIZE)
experiment_group = np.concatenate(
    (
        experiment_group,
        np.random.choice(EXPERIMENT_GROUPS, NUM_CUSTOMERS - len(experiment_group)),
    )
)
customer_group_mapping = dict(zip(customer_ids, experiment_group))

# Generate orders
order_ids = np.arange(1, NUM_ORDERS + 1)
customers = np.random.choice(customer_ids, NUM_ORDERS)
order_values = np.abs(
    np.random.normal(loc=10, scale=2, size=NUM_ORDERS)
)  # Normally distributed around 10 and positive
order_delivery_times = np.abs(
    np.random.normal(loc=30, scale=5, size=NUM_ORDERS)
)  # Normally distributed around 30 minutes and positive
order_city_codes = np.random.randint(
    1, 3, NUM_ORDERS
)  # Random city codes between 1 and 2

# Create DataFrame
data = {
    "order_id": order_ids,
    "customer_id": customers,
    "experiment_group": [
        customer_group_mapping[customer_id] for customer_id in customers
    ],
    "order_value": order_values,
    "order_delivery_time": order_delivery_times,
    "order_city_code": order_city_codes,
}

df = pd.DataFrame(data)
df.order_city_code = df.order_city_code.astype(str)

# -------------------------------------------------
# ---- Define metrics, variants and dimensions ----
# -------------------------------------------------

dimension__city_code = Dimension(name="order_city_code", values=["1", "2"])

metric__order_value = SimpleMetric(alias="AOV", name="order_value")

metric__delivery_time = SimpleMetric(
    alias="AVG DT", name="order_delivery_time_in_minutes"
)

variants = [
    Variant("control", is_control=True),
    Variant("treatment_1", is_control=False),
    Variant("treatment_2", is_control=False),
]

# --------------------------------------------------------
# ---- Define a simple analysis plan from the metrics ----
# --------------------------------------------------------

simple_analysis_plan = AnalysisPlan.from_metrics(
    metrics=[metric__delivery_time, metric__order_value],
    variants=variants,
    variant_col="experiment_group",
    alpha=0.01,
    dimensions=[dimension__city_code],
    analysis_type="clustered_ols",
    analysis_config={"cluster_cols": ["customer_id"]},
)

# --------------------------------------------------------
# ---- Run analysis and get the experiment scorecard  ----
# --------------------------------------------------------

simple_results = simple_analysis_plan.analyze(exp_data=df, verbose=True)

simple_results_df = simple_results.to_dataframe()