# Authenticate to Google Cloud and configure your project.
from google.colab import auth
auth.authenticate_user()

import os
PROJECT_ID = "your-gcp-project-id"   # FILL IN
DATASET_ID = "cs145_p2"              # FILL IN; must exist in PROJECT_ID
LOCATION = "US"

os.environ['GOOGLE_CLOUD_PROJECT'] = PROJECT_ID

from google.cloud import bigquery
bq = bigquery.Client(project=PROJECT_ID, location=LOCATION)
print(f"Using project={PROJECT_ID}, dataset={DATASET_ID}, location={LOCATION}")

def estimate_query_cost(sql: str) -> dict:
    """Dry-run a query; return bytes scanned + USD estimate."""
    job_config = bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)
    job = bq.query(sql, job_config=job_config)
    bytes_processed = job.total_bytes_processed
    # On-demand pricing: $6.25 / TB scanned (us multi-region).
    cost_usd = bytes_processed / 1e12 * 6.25
    return {
        "bytes": bytes_processed,
        "gb": round(bytes_processed / 1e9, 3),
        "cost_usd": round(cost_usd, 4),
    }

# Quick sanity check (free; no charge for dry-run).
example_sql = f"SELECT COUNT(*) FROM `{PROJECT_ID}.{DATASET_ID}.INFORMATION_SCHEMA.TABLES`"
print(estimate_query_cost(example_sql))

# Recipe A: FHIR / healthcare
FHIR_KEYWORDS = {
    "diabetes", "hypertension", "asthma", "copd", "myocardial",
    "stroke", "ischemic", "anemia", "hyperlipidemia", "obesity",
    "depression", "anxiety", "psoriasis", "arthritis", "pneumonia",
    "metformin", "insulin", "statin", "lisinopril", "amlodipine",
    "icd-10", "snomed", "rxnorm", "loinc", "fhir",
    "cardiology", "oncology", "neurology", "pulmonology", "endocrinology",
}

# Recipe B: NYC taxi / urban
NYC_KEYWORDS = {
    "manhattan", "brooklyn", "queens", "the bronx", "staten island",
    "harlem", "soho", "tribeca", "chelsea", "midtown",
    "jfk airport", "laguardia airport", "new york city subway",
    "metropolitan transportation authority", "yellow cab",
    "new york city neighborhood", "new york city street",
}

# Recipe C: Finance / public companies
FINANCE_KEYWORDS = {
    "new york stock exchange", "nasdaq", "s&p 500", "ipo",
    "dividend", "earnings", "merger", "acquisition", "hedge fund",
    "fortune 500", "sec filing", "10-k", "balance sheet",
    "p/e ratio", "market capitalization",
}

# Recipe D: Music
MUSIC_KEYWORDS = {
    "billboard hot 100", "grammy award", "rolling stone",
    "album", "single", "discography", "music genre", "record label",
    "spotify", "apple music", "musician", "band",
}

# Pick one (or define your own keyword set):
MY_KEYWORDS = FHIR_KEYWORDS  # CHANGE THIS
print(f"Using filter with {len(MY_KEYWORDS)} keywords")

!pip install -q pyspark==3.5.0 datasets pyarrow

from pyspark.sql import SparkSession
spark = (
    SparkSession.builder
        .appName("cs145-p2-wiki-elt")
        .config("spark.driver.memory", "8g")
        .config("spark.sql.execution.arrow.pyspark.enabled", "true")
        .getOrCreate()
)
print("Spark version:", spark.version)

# Download parquet files from HuggingFace Hub (no pandas roundtrip).
!pip install -q huggingface_hub

from huggingface_hub import snapshot_download
LOCAL_WIKI_DIR = snapshot_download(
    repo_id="wikimedia/wikipedia",
    allow_patterns=["20231101.en/*.parquet"],
    repo_type="dataset",
)
print("Wikipedia parquet downloaded to:", LOCAL_WIKI_DIR)

# Spark reads the parquet directory in parallel.
from pyspark.sql.functions import udf, col, lower
from pyspark.sql.types import BooleanType, StringType
import re

wiki = spark.read.parquet(f"{LOCAL_WIKI_DIR}/20231101.en")
print("Raw articles in dump:", wiki.count())

# Filter: keep articles whose title or first 2k chars mention a domain keyword.
# (UDFs are slower than SQL `like`; for large keyword sets, consider rewriting
# this as a regex-based filter on lower(title) and lower(substr(text, 1, 2000)).)
def in_domain(title, text):
    if not title or not text: return False
    t = title.lower()
    body = text[:2000].lower()
    return any(k in t or k in body for k in MY_KEYWORDS)

in_domain_udf = udf(in_domain, BooleanType())

def first_paragraph(text):
    if not text: return ""
    parts = re.split(r"\n\n", text, maxsplit=2)
    return (parts[0] if parts else "")[:2000]

first_para_udf = udf(first_paragraph, StringType())

slim = (
    wiki.filter(in_domain_udf(col("title"), col("text")))
        .select(
            col("id").alias("article_id"),
            col("title"),
            col("url"),
            first_para_udf(col("text")).alias("first_paragraph"),
        )
)

# Checkpoint to Colab disk as a single parquet file (coalesce(1)) so the
# downstream BQ load is one upload, not N. For 1 GB this is fine; if you
# scale up at tier 87+, drop coalesce(1) and load each part separately.
CHECKPOINT_DIR = "/content/wiki_domain.parquet"
slim.coalesce(1).write.mode("overwrite").parquet(CHECKPOINT_DIR)
n_articles = spark.read.parquet(CHECKPOINT_DIR).count()
print(f"Kept {n_articles:,} in-domain articles")

TABLE_WIKI_DOMAIN = f"{PROJECT_ID}.{DATASET_ID}.wiki_domain"

# Spark wrote a single part file because we used coalesce(1).
from pathlib import Path
parts = sorted(Path(CHECKPOINT_DIR).glob("part-*.parquet"))
assert len(parts) == 1, f"expected one part file (coalesce(1)); got {len(parts)}"

job_config = bigquery.LoadJobConfig(
    source_format=bigquery.SourceFormat.PARQUET,
    write_disposition="WRITE_TRUNCATE",
)
with open(parts[0], "rb") as f:
    job = bq.load_table_from_file(f, TABLE_WIKI_DOMAIN, job_config=job_config)
job.result()

print(f"Loaded to {TABLE_WIKI_DOMAIN}")
print(bq.get_table(TABLE_WIKI_DOMAIN))

sql = f"""
SELECT article_id, title, SUBSTR(first_paragraph, 1, 200) AS preview
FROM `{TABLE_WIKI_DOMAIN}`
ORDER BY RAND()
LIMIT 30
"""
sample = bq.query(sql).to_dataframe()
sample

sql = f"""
-- FILL IN: replace this with a real join between your fact table and wiki_domain.
-- Example pattern (FHIR):
-- SELECT c.condition_name, COUNT(w.article_id) AS wiki_coverage
-- FROM `{PROJECT_ID}.{DATASET_ID}.fhir_conditions` c
-- LEFT JOIN `{TABLE_WIKI_DOMAIN}` w
--   ON LOWER(w.title) LIKE CONCAT('%', LOWER(c.condition_name), '%')
-- GROUP BY c.condition_name
-- ORDER BY wiki_coverage DESC
SELECT COUNT(*) AS row_count FROM `{TABLE_WIKI_DOMAIN}`
"""

print("Cost preview:", estimate_query_cost(sql))
df_join = bq.query(sql).to_dataframe()
df_join.head(20)

# Compare title-keyword hit counts: my slice vs public sample.
parity_sql = f"""
WITH keyword_hits AS (
  SELECT 'mine'   AS source, COUNT(*) AS hits
  FROM `{TABLE_WIKI_DOMAIN}`
  WHERE LOWER(title) LIKE '%diabetes%'   -- replace with one of YOUR keywords
  UNION ALL
  SELECT 'public' AS source, COUNT(*) AS hits
  FROM `bigquery-public-data.samples.wikipedia`
  WHERE LOWER(title) LIKE '%diabetes%'
)
SELECT * FROM keyword_hits
"""
print("Cost preview:", estimate_query_cost(parity_sql))
df_parity = bq.query(parity_sql).to_dataframe()
df_parity

import matplotlib.pyplot as plt

# FILL IN with your real join result.
df_join.plot(kind="bar", x=df_join.columns[0], y=df_join.columns[1], figsize=(10, 5))
plt.title("Phase 0: your dataset × Wikipedia keyword join")
plt.tight_layout()
plt.show()

# One-time: create a Vertex AI connection. Skip if already done.
# Run this in a separate terminal or via shell:
# !bq mk --connection --display_name="vertex_ai_connection" \
#   --connection_type=CLOUD_RESOURCE --project_id=$PROJECT_ID --location=US vertex_ai_connection

create_model_sql = f"""
CREATE OR REPLACE MODEL `{PROJECT_ID}.{DATASET_ID}.embed_model`
REMOTE WITH CONNECTION `{PROJECT_ID}.us.vertex_ai_connection`
OPTIONS(endpoint = 'text-embedding-005')
"""
bq.query(create_model_sql).result()
print("embed_model created")

embed_sql = f"""
CREATE OR REPLACE TABLE `{PROJECT_ID}.{DATASET_ID}.wiki_domain_embedded` AS
SELECT
  article_id, title, url, first_paragraph,
  ml_generate_embedding_result AS embedding,
  ml_generate_embedding_statistics
FROM ML.GENERATE_EMBEDDING(
  MODEL `{PROJECT_ID}.{DATASET_ID}.embed_model`,
  (SELECT article_id, title, url, first_paragraph,
          first_paragraph AS content
   FROM `{TABLE_WIKI_DOMAIN}`)
)
WHERE ml_generate_embedding_status = ''  -- successful only
"""

print("Cost preview:", estimate_query_cost(embed_sql))
print("Running embedding job. This takes a few minutes for ~50K rows.")
bq.query(embed_sql).result()
print("Done. Embeddings landed in wiki_domain_embedded.")

# Optional, only if your embedded table is ≥ 5 GB:
# CREATE OR REPLACE VECTOR INDEX wiki_idx
#   ON `{PROJECT_ID}.{DATASET_ID}.wiki_domain_embedded`(embedding)
#   OPTIONS(distance_type='COSINE', index_type='IVF')

semantic_sql = f"""
-- VECTOR_SEARCH returns columns: query.*, base.*, distance.
-- The query CTE supplies one or more (query_text, embedding) rows; each is
-- ranked against the base table.
WITH query_embeddings AS (
  SELECT
    query_text,
    ml_generate_embedding_result AS embedding
  FROM ML.GENERATE_EMBEDDING(
    MODEL `{PROJECT_ID}.{DATASET_ID}.embed_model`,
    (
      -- FILL IN: replace with the query strings you actually want to test.
      SELECT 'diabetes type 2' AS content, 'diabetes type 2' AS query_text
      UNION ALL
      SELECT 'metformin dosing',           'metformin dosing'
    )
  )
)
SELECT
  query.query_text,
  base.title,
  base.url,
  distance
FROM VECTOR_SEARCH(
  TABLE `{PROJECT_ID}.{DATASET_ID}.wiki_domain_embedded`,
  'embedding',
  TABLE query_embeddings,
  top_k => 5,
  distance_type => 'COSINE'
)
ORDER BY query.query_text, distance ASC
"""

print("Cost preview:", estimate_query_cost(semantic_sql))
df_sem = bq.query(semantic_sql).to_dataframe()
df_sem

sql = f"""
-- Query 1: FILL IN
SELECT 1
"""
print("Cost preview:", estimate_query_cost(sql))
df_q1 = bq.query(sql).to_dataframe()
df_q1.head(20)

# Visualization for Query 1 (delete this cell if no viz).
import matplotlib.pyplot as plt
# FILL IN: chart df_q1.

sql = f"""
-- Query 2: FILL IN
SELECT 1
"""
print("Cost preview:", estimate_query_cost(sql))
df_q2 = bq.query(sql).to_dataframe()
df_q2.head(20)

# Visualization for Query 2 (delete this cell if no viz).
import matplotlib.pyplot as plt
# FILL IN: chart df_q2.

sql = f"""
-- Query 3: FILL IN
SELECT 1
"""
print("Cost preview:", estimate_query_cost(sql))
df_q3 = bq.query(sql).to_dataframe()
df_q3.head(20)

# Visualization for Query 3 (delete this cell if no viz).
import matplotlib.pyplot as plt
# FILL IN: chart df_q3.

sql = f"""
-- Query 4: FILL IN
SELECT 1
"""
print("Cost preview:", estimate_query_cost(sql))
df_q4 = bq.query(sql).to_dataframe()
df_q4.head(20)

sql = f"""
-- Query 5: FILL IN
SELECT 1
"""
print("Cost preview:", estimate_query_cost(sql))
df_q5 = bq.query(sql).to_dataframe()
df_q5.head(20)

sql = f"""
-- Query 6: FILL IN
SELECT 1
"""
print("Cost preview:", estimate_query_cost(sql))
df_q6 = bq.query(sql).to_dataframe()
df_q6.head(20)

training_sql = f"""
-- FILL IN: build the training table. Include features, target, split column.
SELECT 1 AS feature1, 0 AS target, IF(RAND() < 0.8, 'TRAIN', 'TEST') AS split
"""
print("Cost preview:", estimate_query_cost(training_sql))

train_sql = f"""
CREATE OR REPLACE MODEL `{PROJECT_ID}.{DATASET_ID}.my_model`
OPTIONS(
  model_type = 'linear_reg',                       -- or 'logistic_reg', 'boosted_tree_*'
  input_label_cols = ['target'],
  data_split_col = 'split',
  data_split_method = 'CUSTOM'
) AS
SELECT * FROM ({training_sql})
WHERE split = 'TRAIN'
"""
# bq.query(train_sql).result()

eval_sql = f"""
SELECT * FROM ML.EVALUATE(
  MODEL `{PROJECT_ID}.{DATASET_ID}.my_model`,
  (SELECT * FROM ({training_sql}) WHERE split = 'TEST')
)
"""
# bq.query(eval_sql).to_dataframe()

# Paste one of your complex queries here.
sql_q1_plan = "SELECT 1"
print("Dry-run cost:", estimate_query_cost(sql_q1_plan))
# bq.query(sql_q1_plan).result()

# Paste one of your complex queries here.
sql_q2_plan = "SELECT 1"
print("Dry-run cost:", estimate_query_cost(sql_q2_plan))
# bq.query(sql_q2_plan).result()

# FILL IN: rerun your Spark filter with a broader keyword set or no filter
# at all. Target ≥ 5 GB on disk after BQ load.

layout_sql = f"""
CREATE OR REPLACE TABLE `{PROJECT_ID}.{DATASET_ID}.wiki_domain_partitioned`
PARTITION BY DATE_TRUNC(load_date, MONTH)   -- FILL IN: choose your column
CLUSTER BY category, language               -- FILL IN: 1-4 columns
AS
SELECT * FROM `{PROJECT_ID}.{DATASET_ID}.wiki_domain_large`
"""
# bq.query(layout_sql).result()

mv_sql = f"""
CREATE OR REPLACE MATERIALIZED VIEW `{PROJECT_ID}.{DATASET_ID}.feature_v` AS
SELECT
  -- FILL IN: precomputed aggregate that your ML model uses.
  article_id,
  COUNT(*) AS link_count
FROM `{PROJECT_ID}.{DATASET_ID}.wiki_domain_partitioned`
GROUP BY article_id
"""
# bq.query(mv_sql).result()

before_sql = "-- FILL IN: query against wiki_domain_large (no partition/cluster)"
after_sql  = "-- FILL IN: same query against wiki_domain_partitioned"

print("Before:", estimate_query_cost(before_sql))
print("After: ", estimate_query_cost(after_sql))

scaled_embed_sql = f"""
CREATE OR REPLACE TABLE `{PROJECT_ID}.{DATASET_ID}.wiki_domain_embedded_1gb` AS
SELECT
  article_id, title, url, first_paragraph,
  ml_generate_embedding_result AS embedding
FROM ML.GENERATE_EMBEDDING(
  MODEL `{PROJECT_ID}.{DATASET_ID}.embed_model`,
  (SELECT article_id, title, url, first_paragraph,
          first_paragraph AS content
   FROM `{PROJECT_ID}.{DATASET_ID}.wiki_domain_large`)   -- FILL IN: broader slice
)
WHERE ml_generate_embedding_status = ''
"""
print("Cost preview:", estimate_query_cost(scaled_embed_sql))
print("This will charge $8 to $15. Confirm before running.")
# bq.query(scaled_embed_sql).result()

vec_index_sql = f"""
CREATE OR REPLACE VECTOR INDEX wiki_idx_1gb
ON `{PROJECT_ID}.{DATASET_ID}.wiki_domain_embedded_1gb`(embedding)
OPTIONS(distance_type = 'COSINE', index_type = 'IVF')
"""
# bq.query(vec_index_sql).result()

# Poll until the index is READY:
status_sql = f"""
SELECT index_name, coverage_percentage, last_refresh_time
FROM `{PROJECT_ID}.{DATASET_ID}.INFORMATION_SCHEMA.VECTOR_INDEXES`
WHERE table_name = 'wiki_domain_embedded_1gb'
"""
# bq.query(status_sql).to_dataframe()

hybrid_sql = f"""
-- Hybrid retrieval: keyword candidates + vector candidates, fused by reciprocal rank.
WITH query_embedding AS (
  SELECT ml_generate_embedding_result AS embedding
  FROM ML.GENERATE_EMBEDDING(
    MODEL `{PROJECT_ID}.{DATASET_ID}.embed_model`,
    (SELECT 'diabetes type 2' AS content))
),
vec AS (
  SELECT base.article_id, base.title, distance,
         ROW_NUMBER() OVER (ORDER BY distance) AS vec_rank
  FROM VECTOR_SEARCH(
    TABLE `{PROJECT_ID}.{DATASET_ID}.wiki_domain_embedded_1gb`, 'embedding',
    TABLE query_embedding,
    top_k => 50, distance_type => 'COSINE')
),
kw AS (
  SELECT article_id, title,
         ROW_NUMBER() OVER (ORDER BY LENGTH(title)) AS kw_rank
  FROM `{PROJECT_ID}.{DATASET_ID}.wiki_domain_large`
  WHERE LOWER(title) LIKE '%diabetes%'    -- FILL IN: align with query
  LIMIT 50
)
SELECT
  COALESCE(vec.article_id, kw.article_id) AS article_id,
  COALESCE(vec.title, kw.title) AS title,
  1.0/(60 + COALESCE(vec.vec_rank, 1000)) +
  1.0/(60 + COALESCE(kw.kw_rank, 1000)) AS rrf_score
FROM vec FULL OUTER JOIN kw USING (article_id, title)
ORDER BY rrf_score DESC
LIMIT 10
"""
print("Cost preview:", estimate_query_cost(hybrid_sql))
df_hybrid = bq.query(hybrid_sql).to_dataframe()
df_hybrid

# Brute-force version: same query, but force VECTOR_SEARCH to bypass the index.
brute_sql = hybrid_sql.replace(
    "top_k => 50, distance_type => 'COSINE'",
    "top_k => 50, distance_type => 'COSINE', options => '{\"use_brute_force\": true}'"
)

import time
t0 = time.time(); bq.query(hybrid_sql).result(); t_indexed = time.time() - t0
t0 = time.time(); bq.query(brute_sql).result();  t_brute   = time.time() - t0

print(f"Indexed:     {t_indexed:.2f}s, {estimate_query_cost(hybrid_sql)}")
print(f"Brute-force: {t_brute:.2f}s,   {estimate_query_cost(brute_sql)}")

# FILL IN: define your vocabulary as BigQuery table functions (6 to 8 terms).
# Uses bq, PROJECT_ID, DATASET_ID from the Substrate cell at the top.
DS = f"{PROJECT_ID}.{DATASET_ID}"

bq.query(f"""
CREATE OR REPLACE TABLE FUNCTION `{DS}.listener_of`(artist STRING) AS (
  SELECT DISTINCT l.user_id
  FROM `{DS}.Listens` l JOIN `{DS}.Songs` s USING (song_id)
  WHERE s.artist = artist
)
""").result()

# TODO: active_since(days INT64)   -- pin: UTC calendar days, not rolling 24h; NULL time is not active
# TODO: liked(min_rating FLOAT64)  -- pin: NULL rating is not a like
# ... 6 to 8 terms. Each one is a decision you can defend.

# FILL IN: the manifest the agent puts in the model context.
MANIFEST = """
Compose these BigQuery table functions by name; prefer them over raw SQL.
- listener_of(artist STRING): users who played >= 1 song by that artist.
- active_since(days INT64): users with a listen in the last N UTC calendar days.
- liked(min_rating FLOAT64): (user, song) rated >= min_rating; a NULL rating is not a like.
# ... one line per function ...

Schema: <paste your tables and columns>

Join policy: the safe joins are the foreign keys (Listens.user_id = Users.user_id,
Listens.song_id = Songs.song_id). Never join user_id to song_id (both INT, meaningless).
Never join on rating or listen_time.
"""

# FILL IN: the text-to-SQL agent. It MUST compose your stored table functions.
def agent_sql(question, grounded=True):
    context = MANIFEST if grounded else ""      # the "bare" run gets no vocabulary
    prompt = f"""{context}
Write ONE BigQuery Standard SQL query answering the question.
Compose the table functions above by name where they fit; obey the join policy.
Return only SQL.
Question: {question}"""
    # Option A: BigQuery ML.GENERATE_TEXT (zero setup, uses your $50 credit)
    # Option B: any external LLM API you have a key for
    sql = generate_sql(prompt)                  # TODO: implement (ML.GENERATE_TEXT or an API)
    return sql

def answer(question, grounded=True):
    return bq.query(agent_sql(question, grounded)).result().to_dataframe()

# FILL IN: about 10 questions with confirmed answers.
EVALS = [
    ("users who don't listen to Taylor Swift", {"Pluto"}),
    # ... about 10 total; span your functions; include >= 2 gotcha questions
    #     (NULLs, no-activity rows, dedup).
]

def score(grounded):
    hits = 0
    for q, expected in EVALS:
        try:
            got = set(answer(q, grounded=grounded).iloc[:, 0].astype(str))
        except Exception:
            got = None
        hits += (got == expected)
    return hits / len(EVALS)

print("bare:     ", score(grounded=False))
print("grounded: ", score(grounded=True))   # the lift is your headline result