# Generated by Django 5.0.8 on 2024-10-01 14:43

from collections import defaultdict
from contextlib import suppress

from django.db import migrations, models
from tqdm import tqdm


def remove_duplicates(apps, schema_editor):
    changed_portfolios = defaultdict(set)
    from wbportfolio.models import Portfolio

    AssetPosition = apps.get_model("wbportfolio", "AssetPosition")
    for row in (
        AssetPosition.objects.values("portfolio", "date", "underlying_instrument", "portfolio_created")
        .annotate(c=models.Count("*"))
        .filter(c__gt=1)
    ):
        with suppress(Exception):
            qs = AssetPosition.objects.filter(
                date=row["date"],
                portfolio=row["portfolio"],
                underlying_instrument=row["underlying_instrument"],
                portfolio_created=row["portfolio_created"],
            )
            qs.filter(weighting=0).delete()
            base_p = qs.first()
            for a in qs.exclude(id=base_p.id):
                if (
                    abs(base_p.weighting / a.weighting - 1) < 0.02
                    or (a.import_source is None and a.portfolio_created is None)
                    or (a.portfolio.id == 117 and a.underlying_instrument.id == 214)
                ):
                    changed_portfolios[a.portfolio].add(a.date)
                    a.delete()
            if base_p.import_source and (
                qs.filter(import_source=base_p.import_source).count() > 1
                or (base_p.portfolio.id == 123 and base_p.underlying_instrument.id == 214)
            ):
                base_p.weighting += qs.aggregate(c=models.Sum("weighting"))["c"]
                base_p.shares += qs.aggregate(c=models.Sum("shares"))["c"]
                base_p.save()
                changed_portfolios[base_p.portfolio].add(base_p.date)
                qs.filter(import_source=base_p.import_source).delete()

    # Normalizing portfolio weights
    for portfolio, dates in tqdm(changed_portfolios.items()):
        portfolio = Portfolio.all_objects.get(id=portfolio.id)
        for _d in dates:
            portfolio.change_at_date(_d)


class Migration(migrations.Migration):
    dependencies = [
        ("currency", "0001_initial"),
        ("io", "0007_alter_exportsource_query_params"),
        ("wbfdm", "0026_instrument_is_cash_equivalent"),
        ("wbportfolio", "0066_assetposition_initial_shares_at_custodian_and_more"),
    ]

    operations = [
        migrations.RunSQL(
            sql="SET CONSTRAINTS ALL IMMEDIATE;",
        ),
        migrations.RunPython(remove_duplicates),
        migrations.RunSQL(
            sql="SET CONSTRAINTS ALL DEFERRED;",
        ),
        migrations.AddConstraint(
            model_name="assetposition",
            constraint=models.UniqueConstraint(
                fields=("portfolio", "date", "underlying_instrument", "portfolio_created"),
                name="unique_asset_position",
            ),
        ),
    ]
