import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

np.random.seed(42)

# 8 BC municipalities, 24 months, treatment at month 13
cities = {
    "Nanaimo":       {"base": 1250, "trend": 3.0, "treated": False},
    "Kelowna":       {"base": 1580, "trend": 4.5, "treated": False},
    "Kamloops":      {"base": 1150, "trend": 2.5, "treated": False},
    "Victoria":      {"base": 1620, "trend": 4.0, "treated": False},
    "Prince George": {"base": 980,  "trend": 1.5, "treated": False},
    "Squamish":      {"base": 1720, "trend": 5.0, "treated": True},
    "Tofino":        {"base": 1600, "trend": 6.0, "treated": True},
    "Nelson":        {"base": 1450, "trend": 5.5, "treated": True},
}

rows = []
for city, info in cities.items():
    for month in range(1, 25):
        post = int(month > 12)
        treated = int(info["treated"])
        effect = 0 if not info["treated"] or not post else np.random.uniform(40, 60)
        rent = (info["base"]
                + info["trend"] * month
                - treated * post * effect
                + np.random.normal(0, 15))
        rows.append({
            "municipality": city,
            "month": month,
            "post": post,
            "short_term_rental_ban": treated,
            "avg_rent": round(rent),
        })

df = pd.DataFrame(rows)

print(df.head(16).to_string(index=False))
print(f"\nShape: {df.shape}")

# Aggregate to group-level means by month
treated_avg = df[df.short_term_rental_ban == 1].groupby("month")["avg_rent"].mean()
control_avg = df[df.short_term_rental_ban == 0].groupby("month")["avg_rent"].mean()

# Pre-treatment control trend (slope)
pre_ctrl = control_avg.loc[1:12]
trend = (pre_ctrl.iloc[-1] - pre_ctrl.iloc[0]) / 11

# Counterfactual: treated baseline + control growth from that baseline
treated_baseline = treated_avg.loc[12]
cf_months = range(13, 25)
counterfactual = [treated_baseline + trend * (m - 12) for m in cf_months]

# Post-treatment treated mean
post_tr = treated_avg.loc[13:24].mean()
post_cf = np.mean(counterfactual)
att = post_tr - post_cf

fig, ax = plt.subplots(figsize=(10, 5))

# Plot lines
ax.plot(treated_avg.index, treated_avg.values, 'o-', label="Treated (ban adopted)", color="#2c7fb8", markersize=4)
ax.plot(control_avg.index, control_avg.values, 's-', label="Control (no ban)", color="#bdbdbd", markersize=4)
ax.plot(cf_months, counterfactual, 's--', label="Counterfactual", color="#2c7fb8", alpha=0.5, markersize=4)

# Treatment line
ax.axvline(12.5, color="black", linestyle=":", linewidth=1)
ax.text(12.7, ax.get_ylim()[1] - 130, "Treatment", fontsize=9, color="black")

# ATT bracket
post_x = 18  # mid-post period
ax.annotate("", xy=(post_x, post_tr), xytext=(post_x, post_cf),
            arrowprops=dict(arrowstyle="<->", color="#e34a33", lw=1.5))
ax.text(post_x + 0.3, post_cf - 50, f"ATT ≈ -${-1*att:.0f}/mo",
        fontsize=10, color="#e34a33", va="center")

ax.set_xlabel("Month")
ax.set_ylabel("Average Rent ($)")
ax.legend(loc="upper left")
ax.set_title("Difference-in-Differences: Short-Term Rental Ban and Average Rent")
plt.tight_layout()
plt.show()

print(f"Control trend (pre): ${trend:.1f}/month")
print(f"Treated baseline (month 12): ${treated_baseline:.0f}")
print(f"Post-treatment treated avg: ${post_tr:.0f}")
print(f"Counterfactual avg: ${post_cf:.0f}")
print(f"Estimated ATT: ${att:.0f}/month")

treated_avg.loc[13:24].mean() - control_avg.loc[13:24].mean()

import cvxpy as cp

# --- Pivot the data ---
pivot = df.pivot(index="month", columns="municipality", values="avg_rent")

treated_munis = [c for c, i in cities.items() if i["treated"]]
control_munis = [c for c, i in cities.items() if not i["treated"]]

y_tr = pivot[treated_munis]
y_co = pivot[control_munis]

y_pre_co  = y_co.loc[1:12]
y_pre_tr  = y_tr.loc[1:12]
y_post_co = y_co.loc[13:24]
y_post_tr = y_tr.loc[13:24]

# --- Fit canonical synthetic control ---
y_target = y_pre_tr.mean(axis=1).values
X = y_pre_co.values

w = cp.Variable(X.shape[1])
objective = cp.Minimize(cp.sum_squares(X @ w - y_target))
constraints = [cp.sum(w) == 1, w >= 0]
problem = cp.Problem(objective, constraints)
problem.solve()

weights = w.value
synth_pre  = X @ weights
synth_post = y_post_co.values @ weights

print("Weights (non-zero):")
for name, wt in zip(control_munis, weights):
    if wt > 0.001:
        print(f"  {name}: {wt:.3f}")

# --- Plot ---
from matplotlib.lines import Line2D
fig, ax = plt.subplots(figsize=(10, 5))

for muni in control_munis:
    wt = weights[control_munis.index(muni)]  # get this city's weight
    color = "gray" if wt < 0.001 else "#1b9e77"  # dark teal for selected
    alpha = 0.25
    ax.plot(pivot.index, pivot[muni], color=color, alpha=alpha, linewidth=0.8)

# Treated average
treated_avg = y_tr.mean(axis=1)
ax.plot(treated_avg.index, treated_avg.values, 'o-', color="#2c7fb8", linewidth=1.8, markersize=4, label="Treated (average)")

# Synthetic control
synth_full = np.concatenate([synth_pre, synth_post])
ax.plot(range(1, 25), synth_full, '--', color="orange", linewidth=1.8, label="Synthetic control")

# Treatment line
ax.axvline(12.5, color="black", linestyle=":", linewidth=1)
ax.text(12.7, ax.get_ylim()[1] - 200, "Treatment", fontsize=9, color="black")

# ATT annotation
post_tr_avg = y_post_tr.mean(axis=1).mean()
post_sc_avg = synth_post.mean()
att = post_tr_avg - post_sc_avg

post_x = 18
treated_y = y_post_tr.mean(axis=1).iloc[5]
synth_y = synth_post[5]
ax.annotate("", xy=(post_x, synth_y-att/2.0), xytext=(post_x, treated_y+att/2.0),
            arrowprops=dict(arrowstyle="<->", color="#e34a33", lw=1.5))
ax.text(post_x + 0.3, post_sc_avg - 80, f"ATT ≈ -${-1*att:.0f}/mo", fontsize=10, color="#e34a33", va="center")

ax.set_xlabel("Month")
ax.set_ylabel("Average Rent ($)")

handles, labels = ax.get_legend_handles_labels()
handles.append(Line2D([0], [0], color="#1b9e77", alpha=0.25, linewidth=0.8))
labels.append("Selected controls")
handles.append(Line2D([0], [0], color="gray", alpha=0.25, linewidth=0.8))
labels.append("Unselected controls")
ax.legend(handles, labels, loc="right")

ax.set_title("Synthetic Control: Short-Term Rental Ban and Average Rent")
plt.tight_layout()
plt.show()

print(f"\nTreated post avg: ${post_tr_avg:.0f}")
print(f"Synthetic control post avg: ${post_sc_avg:.0f}")
print(f"Estimated ATT: ${att:.0f}/month")

list(zip(control_munis, weights))

import statsmodels.formula.api as smf

T_pre = len(y_pre_co)
T_post = len(y_post_co)
T = T_pre + T_post

from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.utils.validation import check_X_y, check_array, check_is_fitted

class SyntheticControl(BaseEstimator, RegressorMixin):

    def __init__(self,):
        pass

    def fit(self, y_pre_co, y_pre_tr):

        y_pre_co, y_pre_tr = check_X_y(y_pre_co, y_pre_tr)
    
        w = cp.Variable(y_pre_co.shape[1])
        
        objective = cp.Minimize(cp.sum_squares(y_pre_co@w - y_pre_tr))
        constraints = [cp.sum(w) == 1, w >= 0]
        
        problem = cp.Problem(objective, constraints)
        
        self.loss_ = problem.solve(verbose=False)
        self.w_ = w.value
        
        self.is_fitted_ = True
        return self
        
    def predict(self, y_co):

        check_is_fitted(self)
        y_co = check_array(y_co)
        
        return y_co @ self.w_

sc_model = SyntheticControl()
sc_model.fit(y_pre_co, y_pre_tr.mean(axis=1))

(y_post_tr.mean(axis=1) - sc_model.predict(y_post_co)).mean()

unit_w = pd.DataFrame(zip(y_pre_co.columns, sc_model.w_),
                         columns=["municipality", "unit_weight"])

unit_w.head()

df.head()

df_with_w = (df
             .assign(tr_post = lambda d: d["post"]*d["short_term_rental_ban"])
             .merge(unit_w, on=["municipality"], how="left")
             .fillna({"unit_weight": df["short_term_rental_ban"].mean()}))
             
df_with_w.head()

mod = smf.wls(
    "avg_rent ~ tr_post + C(month)",
    data=df_with_w.query("unit_weight>=1e-6"),
    weights=df_with_w.query("unit_weight>=1e-6")["unit_weight"]
)

mod.fit().params["tr_post"]

class SyntheticControl(BaseEstimator, RegressorMixin):

    def __init__(self, fit_intercept=False):
        self.fit_intercept = fit_intercept

    def fit(self, y_pre_co, y_pre_tr):

        y_pre_co, y_pre_tr = check_X_y(y_pre_co, y_pre_tr)
        
        # add intercept
        intercept = np.ones((y_pre_co.shape[0], 1))*self.fit_intercept
        X = np.concatenate([intercept, y_pre_co], axis=1)
        w = cp.Variable(X.shape[1])
        
        objective = cp.Minimize(cp.sum_squares(X@w - y_pre_tr))
        constraints = [cp.sum(w[1:]) == 1, w[1:] >= 0]
        
        problem = cp.Problem(objective, constraints)
        
        self.loss_ = problem.solve(verbose=False)
        self.w_ = w.value[1:] # remove intercept
        
        self.is_fitted_ = True
        return self
        
        
    def predict(self, y_co):

        check_is_fitted(self)
        y_co = check_array(y_co)
        
        return y_co @ self.w_

time_sc = SyntheticControl(fit_intercept=True)

time_sc.fit(
    y_pre_co.T,
    y_post_co.mean(axis=0)
)

time_w = pd.DataFrame(zip(y_pre_co.index, time_sc.w_),
                         columns=["month", "time_weight"])

time_w.tail()

plt.figure(figsize=(10,4))
plt.bar(time_w["month"], time_w["time_weight"])
plt.ylabel("Time Weights")
plt.xticks(rotation=45);

scdid_df = (
    df_with_w
    .merge(time_w, on=["month"], how="left")
    .fillna({"time_weight":df["post"].mean()})
    .assign(weight = lambda d: (d["time_weight"]*d["unit_weight"]))
)
scdid_df.head()

did_model = smf.wls(
    "avg_rent ~ short_term_rental_ban*post",
    data=scdid_df.query("weight>1e-10"),
    weights=scdid_df.query("weight>1e-10")["weight"]).fit()

did_model.params["short_term_rental_ban:post"]

tr_period = 12
avg_pre_period = time_w.iloc[int((time_w["time_weight"] * np.arange(0, len(time_w), 1)).sum().round())]["month"]
avg_post_period = 18

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15,8), sharex=True, gridspec_kw={'height_ratios': [3, 1]})

ax1.plot(y_co.index, y_tr.mean(axis=1), label="Treated")
ax1.plot(y_co.index, sc_model.predict(y_co), label="Synthetic Control", ls="-.")
ax1.axvline(tr_period, color="black", ls="dotted")

pre_sc = did_model.params["Intercept"]
post_sc = pre_sc + did_model.params["post"]
pre_treat = pre_sc + did_model.params["short_term_rental_ban"]
post_treat = post_sc + did_model.params["short_term_rental_ban"] + did_model.params["short_term_rental_ban:post"]

sc_did_y0 = pre_treat + (post_sc - pre_sc)

for muni_idx in unit_w.index:
    wt = weights[muni_idx]  # get this city's weight
    color = "gray" if wt < 0.001 else "#1b9e77"  # dark teal for selected
    alpha = 0.25 #if wt < 0.001 else 0.9
    ax1.plot(pivot.index, pivot[control_munis[muni_idx]], color=color, alpha=alpha, linewidth=0.8)

handles, labels = ax.get_legend_handles_labels()
handles.append(Line2D([0], [0], color="#1b9e77", alpha=0.25, linewidth=0.8))
labels.append("Selected controls")
handles.append(Line2D([0], [0], color="gray", alpha=0.25, linewidth=0.8))
labels.append("Unselected controls")
ax1.legend(handles, labels, loc="right")
ax1.set_title("Synthetic Difference-in-Differences")
ax1.set_ylabel("Average Rent")

post_x = 18
appx_att = -1*did_model.params["short_term_rental_ban:post"]
ax1.text(post_x + 0.3, post_sc_avg - 80, f"ATT ≈ -${appx_att:.0f}/mo", fontsize=10, color="#e34a33", va="center")
ax1.annotate("", xy=(post_x, y_post_tr.iloc[5].mean()-appx_att/2.0), xytext=(post_x, synth_post[5]+appx_att/2.0),
            arrowprops=dict(arrowstyle="<->", color="#e34a33", lw=1.5))

ax2.bar(time_w["month"], time_w["time_weight"])
ax2.axvline(tr_period, color="black", ls="dotted")
ax2.set_ylabel("Time Weights")
ax2.set_xlabel("Months")
plt.xticks(rotation=30);

y_post_tr.iloc[5].mean()