from enum import unique
from itertools import groupby
from unicodedata import category

from vespa.package import Field, RankProfile, MatchPhaseRanking
%load_ext autoreload
%autoreload 2
import mycode.vap as vap

# Index about 10k docs
# have 10 categories
# Score random in the first phase
# match-phase-should take top1 group
# search for top-10 hits
#

app = vap.demo_application_package()

from vespa.package import Field, Diversity, MatchPhaseRanking, RankProfile

app.get_schema().document.add_fields(
    Field(name="category", type="string", indexing="attribute | summary", attribute=["fast-search"])
)
app.get_schema().add_rank_profile(
    RankProfile(
        name="matchphase",
        match_phase=MatchPhaseRanking(
            attribute="id",
            order="descending",
            max_hits=10
        ),
        first_phase="attribute(id)",
    ),
)
app.get_schema().add_rank_profile(
    RankProfile(
        name="matchphasediversity",
        match_phase=MatchPhaseRanking(
            attribute="id",
            order="descending",
            max_hits=10
        ),
        diversity=Diversity(
            attribute="category",
            min_groups=10,
        ),
        first_phase="random",
    ),
)

print(app.get_schema().schema_to_text)

from vespa.deployment import VespaDocker

# In case running colima on macos run the following
# !sudo ln -sf $HOME/.colima/default/docker.sock /var/run/docker.sock
vespa_docker = VespaDocker(
    container_image="vespaengine/vespa:8.650.16",
)
# Start a docker container and deploy the application package
client = vespa_docker.deploy(
    application_package=app,
)

vap.feed(client, docs=[{'id': index, 'category': int(index / 1000)} for index in range(10000)])

resp1 = client.query(body={
    'yql': 'select * from sources * where true',
    'hits': 400,
    'ranking.profile': 'matchphasediversity',
}).json
resp1['root']['children']

resp1

from collections import Counter
Counter(list(map(lambda x: x['fields']['category'], resp1['root']['children'])))

resp = client.query(body={
    'yql': 'select * from sources * where true',
    'hits': 400,
    'ranking.profile': 'matchphase',
    # 'ranking.matchPhase.diversity.minGroups': 0,
    # 'ranking.matchPhase.diversity.attribute': "id",
}).json

resp

Counter(list(map(lambda x: x['fields']['category'], resp['root']['children'])))