"""Converters for ordinal datasets (primary type: score).

Options are levels lowest -> highest (key = level index). The number of levels and their
wording vary per question (e.g. 5 stars collapsed into 3 levels), and a minority of
questions use the reversed polarity ("How negative ...?" with very positive -> very
negative) so the model has to read the level texts rather than assume "high = good".
"""
from __future__ import annotations

from typing import Iterator

from jevlike.data.common import (
    GENERIC_SCORE, Example, Scale, balanced_order, clean, cuts, dedupe_rows, make_score, mapping, noul,
    pair_state, reverse, weighted,
)
from jevlike.data.hub import read_many, read_rows
from jevlike.data.sources_choice import Ctx, _tweet_clean


def _score_loop(ctx: Ctx, rows: list, text_fn, value_fn, scales: list[Scale], statements: list,
                *, alpha: float = 0.5, bucket_fn=None, p_noul: float = 0.4, p_second: float = 0.15,
                key_fn=None, **fmt) -> Iterator[Example]:
    """rows -> score question (+ derived statement, + second scale variant).

    statements: list of (text, fn(value) -> True/False/None)."""
    rng = ctx.rng
    for r in balanced_order(rng, rows, bucket_fn or value_fn, alpha):
        v = value_fn(r)
        state = text_fn(r)
        if not state:
            continue
        qs = []
        q = make_score(rng, scales, v, GENERIC_SCORE, **fmt)
        if q:
            qs.append(q)
        if statements and (not qs or rng.random() < p_noul):
            for _ in range(4):
                s, f = rng.choice(statements)
                t = f(v)
                if t is not None:
                    qs.append(noul(s.format(**fmt), t))
                    break
        if qs and rng.random() < p_second:
            q2 = make_score(rng, scales, v, GENERIC_SCORE, **fmt)
            if q2 and q2[0].instructions != qs[0][0].instructions:
                qs.append(q2)
        if qs:
            yield Example(state, qs[:3], key=key_fn(r) if key_fn else None)


# ---------------------------------------------------------------- 1-5 star reviews (value 0..4)

def star_scales(thing: str = "business") -> list[Scale]:
    five = cuts(1, 2, 3, 4)
    S = [
        Scale(["1 star", "2 stars", "3 stars", "4 stars", "5 stars"], five,
              ["How many stars did the reviewer give?", "Predict the star rating of this review.",
               "Rate this review from 1 to 5 stars.", "Star rating?", "What rating (1-5 stars) goes with this review?"],
              1.4),
        Scale(["very negative", "negative", "neutral or mixed", "positive", "very positive"], five,
              ["How positive is this review?", "What is the overall sentiment of the review?", "Rate the sentiment.",
               "Sentiment of the customer?"], 1.2),
        Scale(["terrible", "poor", "average", "good", "excellent"], five,
              ["How would the reviewer rate the experience?", "Overall rating?",
               "What did the customer think of the {thing}?", "How good was the {thing} according to the review?"]),
        Scale(["the reviewer hated it", "the reviewer was mostly dissatisfied", "the reviewer had mixed feelings",
               "the reviewer was mostly satisfied", "the reviewer loved it"], five,
              ["How satisfied is the customer?", "How happy was the reviewer?"], 0.7),
        Scale(["very unlikely", "unlikely", "not sure", "likely", "very likely"], five,
              ["How likely is this customer to recommend the {thing}?", "Would the reviewer come back?"], 0.4),
        Scale(["negative", "mixed", "positive"], mapping({0: 0, 1: 0, 2: 1, 3: 2, 4: 2}),
              ["Is the review negative, mixed or positive?", "Overall sentiment?", "How does the customer feel?"], 1.0),
        Scale(["bad (1-2 stars)", "average (3 stars)", "good (4-5 stars)"], mapping({0: 0, 1: 0, 2: 1, 3: 2, 4: 2}),
              ["Bucket this review by rating.", "Which rating band fits the review?"], 0.5),
        Scale(["unhappy", "neutral", "happy"], mapping({0: 0, 1: 0, 2: 1, 3: 2, 4: 2}),
              ["How happy is the customer?", "Customer mood?"], 0.5),
        Scale(["dissatisfied", "satisfied"], mapping({0: 0, 1: 0, 3: 1, 4: 1}),
              ["Was the customer satisfied?", "Satisfaction?"], 0.4),
        Scale(["negative", "mixed", "good", "excellent"], mapping({0: 0, 1: 0, 2: 1, 3: 2, 4: 3}),
              ["How positive is the review?", "Rate the review."], 0.4),
    ]
    rev5 = reverse(S[1], ["How negative is this review?", "How dissatisfied is the customer?",
                          "How much did the reviewer dislike the {thing}?", "Rate the negativity of the review."], weight=0.7)
    rev3 = reverse(S[7], ["How unhappy is the customer?", "Rate the level of customer dissatisfaction."], weight=0.4)
    return S + [rev5, rev3]


STAR_STATEMENTS = [
    ("The review is positive.", lambda v: True if v >= 3 else (False if v <= 1 else None)),
    ("The review is negative.", lambda v: True if v <= 1 else (False if v >= 3 else None)),
    ("The customer was satisfied with the {thing}.", lambda v: True if v >= 3 else (False if v <= 1 else None)),
    ("The reviewer recommends the {thing}.", lambda v: True if v >= 3 else (False if v <= 1 else None)),
    ("The customer is unhappy.", lambda v: True if v <= 1 else (False if v >= 3 else None)),
    ("The reviewer gave 5 stars.", lambda v: v == 4),
    ("The reviewer gave only 1 star.", lambda v: v == 0),
    ("The rating is at least 4 stars.", lambda v: v >= 3),
    ("The rating is 3 stars or lower.", lambda v: v <= 2),
    ("Is this a positive review?", lambda v: True if v >= 3 else (False if v <= 1 else None)),
]


def yelp(ctx: Ctx):
    rows = read_rows("Yelp/yelp_review_full", "yelp_review_full/test-00000-of-00001.parquet")
    rows = dedupe_rows(rows, lambda r: r["text"], lambda r: r["label"])
    yield from _score_loop(ctx, rows, lambda r: clean(r["text"]), lambda r: int(r["label"]), star_scales(),
                           STAR_STATEMENTS, alpha=1.0, thing="business")


def amazon(ctx: Ctx):
    rows = read_many("SetFit/amazon_reviews_multi_en", ["test.jsonl", "validation.jsonl"])
    rows = dedupe_rows(rows, lambda r: r["text"], lambda r: r["label"])
    yield from _score_loop(ctx, rows, lambda r: clean(r["text"]), lambda r: int(r["label"]), star_scales("product"),
                           STAR_STATEMENTS, alpha=1.0, thing="product")


# ---------------------------------------------------------------- sentiment scales

def sst5(ctx: Ctx):
    rows = read_many("SetFit/sst5", ["train.jsonl", "dev.jsonl", "test.jsonl"])
    rows = dedupe_rows(rows, lambda r: r["text"], lambda r: r["label"])
    five = cuts(1, 2, 3, 4)
    three = mapping({0: 0, 1: 0, 2: 1, 3: 2, 4: 2})
    scales = [
        Scale(["very negative", "negative", "neutral", "positive", "very positive"], five,
              ["What is the sentiment of this movie review snippet?", "How positive is this text?",
               "Rate the sentiment.", "Sentiment?"], 1.5),
        Scale(["hated it", "disliked it", "mixed feelings", "liked it", "loved it"], five,
              ["How did the critic feel about the movie?", "What is the reviewer's opinion of the film?"]),
        Scale(["1/5", "2/5", "3/5", "4/5", "5/5"], five, ["What score would the critic give?",
                                                           "Predict the critic's rating out of 5."], 0.5),
        Scale(["negative", "neutral", "positive"], three, ["Is this negative, neutral or positive?",
                                                           "Overall polarity?"], 1.0),
        Scale(["negative", "positive"], mapping({0: 0, 1: 0, 3: 1, 4: 1}),
              ["Is the opinion negative or positive?"], 0.4),
    ]
    scales.append(reverse(scales[0], ["How negative is this text?", "Rate how critical this review is."], weight=0.6))
    stm = [("The text expresses a positive opinion.", lambda v: True if v >= 3 else (False if v <= 1 else None)),
           ("The reviewer liked the movie.", lambda v: True if v >= 3 else (False if v <= 1 else None)),
           ("The critic is negative about the film.", lambda v: True if v <= 1 else (False if v >= 3 else None)),
           ("This is a strongly worded opinion.", lambda v: True if v in (0, 4) else (False if v == 2 else None)),
           ("The sentiment is neutral.", lambda v: v == 2)]
    yield from _score_loop(ctx, rows, lambda r: clean(r["text"]), lambda r: int(r["label"]), scales, stm, alpha=0.7)


def tweet_sentiment(ctx: Ctx):
    rows = read_many("cardiffnlp/tweet_eval", [f"sentiment/{s}-00000-of-00001.parquet" for s in ("train", "validation", "test")])
    rows = dedupe_rows(rows, lambda r: r["text"], lambda r: r["label"])
    ident = mapping({0: 0, 1: 1, 2: 2})
    scales = [
        Scale(["negative", "neutral", "positive"], ident, ["What is the sentiment of this tweet?", "Tweet sentiment?",
                                                           "How positive is the tweet?", "Rate the tone of the tweet."], 1.5),
        Scale(["unhappy or critical", "neutral or factual", "happy or supportive"], ident,
              ["How does the author feel?", "What is the emotional tone of this post?"]),
        Scale(["negative", "positive"], mapping({0: 0, 2: 1}), ["Is the tweet negative or positive?"], 0.4),
    ]
    scales.append(reverse(scales[0], ["How negative is the tweet?", "Rate the negativity of this post."], weight=0.5))
    stm = [("The tweet is positive.", lambda v: v == 2), ("The tweet expresses a negative sentiment.", lambda v: v == 0),
           ("The tweet is neutral in tone.", lambda v: v == 1),
           ("The author sounds upset or critical.", lambda v: True if v == 0 else (False if v == 2 else None))]
    yield from _score_loop(ctx, rows, lambda r: _tweet_clean(r["text"]), lambda r: int(r["label"]), scales, stm,
                           alpha=0.5)


def fin_sentiment(ctx: Ctx):
    rows = read_many("zeroshot/twitter-financial-news-sentiment", ["sent_train.csv", "sent_valid.csv"])
    order = {0: 0, 2: 1, 1: 2}  # bearish, neutral, bullish
    rows = dedupe_rows(rows, lambda r: r["text"], lambda r: r["label"])
    ident = mapping({0: 0, 1: 1, 2: 2})
    scales = [
        Scale(["bearish", "neutral", "bullish"], ident, ["What is the market sentiment of this tweet?",
                                                         "Is the news bearish or bullish?", "Market sentiment?"], 1.5),
        Scale(["negative for the stock or market", "no clear impact", "positive for the stock or market"], ident,
              ["How does this news affect the stock?", "What is the likely market reaction?"]),
        Scale(["bearish", "bullish"], mapping({0: 0, 2: 1}), ["Bearish or bullish?"], 0.3),
    ]
    scales.append(reverse(scales[0], ["How bearish is this?", "Rate the downside signal in this news."], weight=0.5))
    stm = [("The tweet is bullish.", lambda v: v == 2), ("This news is bad for investors.", lambda v: True if v == 0 else (False if v == 2 else None)),
           ("The tweet is neutral about the market.", lambda v: v == 1)]
    yield from _score_loop(ctx, rows, lambda r: clean(r["text"]), lambda r: order[int(r["label"])], scales, stm, alpha=0.6)


# ---------------------------------------------------------------- similarity (pair states)

STS_INSTR = ["How similar are the two sentences in meaning?", "Rate the semantic similarity of the sentences.",
             "To what extent do the two sentences say the same thing?", "Semantic similarity?",
             "How close in meaning are these texts?"]


def stsb(ctx: Ctx):
    rows = read_many("sentence-transformers/stsb", [f"data/{s}-00000-of-00001.parquet" for s in ("train", "validation", "test")])
    rng = ctx.rng
    scales = [
        Scale(["completely different", "different, but on the same topic", "not equivalent, but share some details",
               "roughly equivalent, some important details differ", "mostly equivalent, only minor details differ",
               "completely equivalent"], lambda v: int(round(v)), STS_INSTR, 1.5),
        Scale(["0", "1", "2", "3", "4", "5"], lambda v: int(round(v)),
              ["Rate the similarity from 0 (unrelated) to 5 (same meaning).", "STS score (0-5)?"], 0.6),
        Scale(["unrelated", "slightly similar", "somewhat similar", "very similar", "equivalent"], cuts(1, 2, 3, 4.2),
              STS_INSTR),
        Scale(["different meaning", "partly similar", "same meaning"], cuts(1.5, 3.8), STS_INSTR + [
            "Do the sentences mean the same thing?"]),
    ]
    scales.append(reverse(scales[2], ["How different are the two sentences in meaning?",
                                      "Rate how much the sentences disagree in meaning."], weight=0.5))
    stm = [("The two sentences mean the same thing.", lambda v: True if v >= 4 else (False if v <= 1.5 else None)),
           ("The sentences are about different things.", lambda v: True if v <= 1 else (False if v >= 3 else None)),
           ("The second sentence paraphrases the first.", lambda v: True if v >= 4.2 else (False if v <= 2 else None))]
    rows = [dict(r, _v=float(r["score"]) * 5) for r in rows]
    rows = [dict(r, _s=pair_state(rng, clean(r["sentence1"]), clean(r["sentence2"]))) for r in rows]
    yield from _score_loop(ctx, rows, lambda r: r["_s"], lambda r: r["_v"], scales, stm, alpha=0.5,
                           bucket_fn=lambda r: int(round(r["_v"])), key_fn=_pair_key)


def _pair_key(r):
    from jevlike.data.common import norm_key
    return " || ".join(sorted([norm_key(str(r["sentence1"])), norm_key(str(r["sentence2"]))]))


def sick(ctx: Ctx):
    rows = read_rows("mteb/sickr-sts", "test.jsonl.gz")
    rng = ctx.rng
    rel = ["How related are these two sentences?", "Rate the semantic relatedness of the pair.",
           "How closely do the two sentences describe the same scene?", "Relatedness?"]
    scales = [
        Scale(["unrelated", "slightly related", "somewhat related", "closely related", "nearly identical meaning"],
              lambda v: min(4, max(0, int(round(v)) - 1)), rel, 1.5),
        Scale(["1", "2", "3", "4", "5"], lambda v: min(4, max(0, int(round(v)) - 1)),
              ["Relatedness score from 1 (unrelated) to 5 (same meaning)?"], 0.6),
        Scale(["low relatedness", "medium relatedness", "high relatedness"], cuts(2.5, 3.8), rel),
    ]
    stm = [("The two sentences describe the same situation.", lambda v: True if v >= 4.5 else (False if v <= 2.5 else None)),
           ("The sentences are unrelated.", lambda v: True if v <= 1.8 else (False if v >= 3.5 else None))]
    rows = [dict(r, _v=float(r["score"]),
                 _s=pair_state(rng, clean(r["sentence1"]), clean(r["sentence2"]))) for r in rows]
    yield from _score_loop(ctx, rows, lambda r: r["_s"], lambda r: r["_v"], scales, stm, alpha=0.4,
                           bucket_fn=lambda r: int(round(r["_v"])), key_fn=_pair_key)


# ---------------------------------------------------------------- toxicity (fraction of raters)

def civil_comments(ctx: Ctx):
    rows = read_rows("google/civil_comments", "data/test-00000-of-00001.parquet")
    rng = ctx.rng
    rows = [r for r in rows if 20 <= len(str(r["text"])) <= 1500]
    rows = dedupe_rows(rows, lambda r: r["text"], lambda r: round(float(r["toxicity"]), 1))

    def dim_scales(dim: str, adj: str, noun: str) -> list[Scale]:
        ins = [f"How {adj} is this comment?", f"Rate the {noun} of the comment.", f"{noun.capitalize()} level?",
               f"How much {noun} does this comment contain?"]
        return [
            Scale([f"not {adj}", f"slightly {adj} or borderline", f"{adj}", f"very {adj}"], cuts(0.1, 0.4, 0.7), ins, 1.0),
            Scale([f"not {adj}", f"borderline", f"clearly {adj}"], cuts(0.15, 0.5), ins, 1.0),
            Scale(["none", "low", "medium", "high"], cuts(0.1, 0.35, 0.65), ins, 0.6),
        ]
    dims = {
        "toxicity": (dim_scales("toxic", "toxic", "toxicity") + [
            Scale(["very toxic", "toxic", "borderline", "civil"], lambda v: 3 - cuts(0.1, 0.4, 0.7)(v),
                  ["How civil is this comment?", "Rate the civility of this comment.",
                   "How respectful is the comment?"], 0.8)], 3.0),
        "insult": (dim_scales("insulting", "insulting", "insult"), 1.0),
        "obscene": (dim_scales("obscene", "vulgar", "obscenity"), 0.6),
        "identity_attack": (dim_scales("identity-attacking", "hostile toward an identity group", "identity attack")[:2], 0.4),
    }
    stm = [("This comment is toxic.", "toxicity", 0.5, 0.1), ("The comment contains an insult.", "insult", 0.5, 0.1),
           ("The comment uses obscene or vulgar language.", "obscene", 0.5, 0.05),
           ("The comment contains a threat.", "threat", 0.5, 0.05),
           ("The comment attacks people for their identity.", "identity_attack", 0.5, 0.05),
           ("This comment is civil and respectful.", "toxicity", -0.1, -0.5)]
    for r in balanced_order(rng, rows, lambda r: min(3, cuts(0.1, 0.4, 0.7)(float(r["toxicity"]))), 0.25):
        qs = []
        dim = weighted(rng, list(dims), [w for _, w in dims.values()])
        q = make_score(rng, dims[dim][0], float(r[dim]))
        if q:
            qs.append(q)
        if rng.random() < 0.5:
            for _ in range(4):
                s, d, hi, lo = rng.choice(stm)
                v = float(r[d])
                if hi < 0:  # inverted statement: true when clearly non-toxic
                    t = True if v <= -hi else (False if v >= -lo else None)
                else:
                    t = True if v >= hi else (False if v <= lo else None)
                if t is not None:
                    qs.append(noul(s, t))
                    break
        if qs:
            yield Example(clean(r["text"]), qs)


SOURCES = {
    "yelp_review_full": yelp, "amazon_reviews": amazon, "sst5": sst5, "tweet_eval_sentiment": tweet_sentiment,
    "fin_news_sentiment": fin_sentiment, "stsb": stsb, "sick": sick, "civil_comments": civil_comments,
}