Land use prediction

import os
from xgboost import XGBClassifier
from lightgbm import LGBMClassifier
from sklearn.ensemble import RandomForestClassifier, HistGradientBoostingClassifier

from blocksnet.analysis.land_use.prediction import SpatialClassifier
from blocksnet.machine_learning.strategy.sklearn.ensemble.voting.classification_strategy import SKLearnVotingClassificationStrategy

import warnings
warnings.filterwarnings("ignore")

import pickle
import pandas as pd
import geopandas as gpd
from pathlib import Path

def load_gdfs(root: str | Path, pattern: str = "*.pkl", target_crs: str | None = None) -> list[gpd.GeoDataFrame]:
    """
    Загружает GeoDataFrame'ы из .pkl файлов:
      - если root указывает на один .pkl файл — читаем его и возвращаем список [GDF]
      - если root — папка — рекурсивно ищем по маске pattern и возвращаем список GDF

    Пытаемся заполнить столбцы 'city' и 'country' из структуры путей:
    .../<country>/<city>/<file>.pkl

    Args:
        root: путь к файлу .pkl или папке
        pattern: маска файлов (по умолчанию "*.pkl") — используется только если root папка
        target_crs: если указан, приводим все геоданные к этому CRS

    Returns:
        list[GeoDataFrame]
    """
    p = Path(root)

    def _ensure_gdf(obj):
        """Преобразует объект в GeoDataFrame, если это DataFrame с колонкой 'geometry'."""
        if isinstance(obj, gpd.GeoDataFrame):
            return obj
        if isinstance(obj, pd.DataFrame) and "geometry" in obj.columns:
            return gpd.GeoDataFrame(obj, geometry="geometry", crs=getattr(obj, "crs", None))
        return None

    def _infer_city_country(fp: Path):
        # ожидаем .../<country>/<city>/<file>.pkl
        city = fp.name[:-4] if fp else None
        country = fp.parent.name if len(fp.parents) >= 2 else None
        return city, country

    # Соберём список файлов
    if p.is_file() and p.suffix.lower() == ".pkl":
        files = [p]
    elif p.is_dir():
        files = sorted(p.rglob(pattern))
        files = [f for f in files if f.suffix.lower() == ".pkl"]
    else:
        return []

    gdfs: list[gpd.GeoDataFrame] = []

    for fp in files:
        try:
            with open(fp, "rb") as f:
                obj = pickle.load(f)
        except Exception as e:
            print(f"Пропускаю {fp}: {e}")
            continue

        # Если GeoDataFrame
        gdf = _ensure_gdf(obj)

        # Если dict с GeoDataFrame
        if gdf is None and isinstance(obj, dict):
            for k, v in obj.items():
                gi = _ensure_gdf(v)
                if gi is None:
                    continue
                gi = gi.copy()
                if "city" not in gi.columns or gi["city"].isna().all():
                    gi["city"] = str(k)
                _, country = _infer_city_country(fp)
                if ("country" not in gi.columns) and country:
                    gi["country"] = country
                if target_crs:
                    gi = gi.to_crs(target_crs) if gi.crs else gi.set_crs(target_crs)
                gdfs.append(gi)
            continue

        if gdf is None:
            print(f"Пропускаю {fp}: объект не GeoDataFrame и не dict с GeoDataFrame.")
            continue

        gdf = gdf.copy()
        city, country = _infer_city_country(fp)
        if "city" not in gdf.columns:
            gdf["city"] = city or fp.stem
        if ("country" not in gdf.columns) and country:
            gdf["country"] = country

        if target_crs:
            gdf = gdf.to_crs(target_crs) if gdf.crs else gdf.set_crs(target_crs)

        gdfs.append(gdf)

    return gdfs

MERGE_DICT = {
    'LandUse.RECREATION': 'non_urban',
    'LandUse.SPECIAL': 'non_urban',
    'LandUse.AGRICULTURE': 'non_urban',
    'LandUse.BUSINESS': 'urban',
    'LandUse.RESIDENTIAL': 'urban',
    'LandUse.INDUSTRIAL': 'industrial',
    'LandUse.TRANSPORT': None,
}

russia = load_gdfs('data/blocks/Russia/')

# 1. Инициализация и обучение
BASE_PARAMS = {"random_state": 42, "n_jobs": -1}
CPU = max(1, min(8, os.cpu_count() or 1))
MODEL_PARAMS = {
    "rf": {
        "n_estimators": 120,          # было 200
        "max_depth": 7,
        "class_weight": "balanced",
        "max_samples": 0.25,          # 🔴 бэггинг на подвыборке
        "min_samples_leaf": 10,       # стабилизация и меньше узлов
        **BASE_PARAMS
    },
    "xgb": {
        "n_estimators": 150,          # меньше
        "max_depth": 7,
        "learning_rate": 0.05,
        "subsample": 0.8,             # стахастичность
        "colsample_bytree": 0.8,
        "tree_method": "hist",        # память/скорость
        "n_jobs": CPU                 # XGB игнорирует BASE_PARAMS если его стерли
    },
    "lgb": {
        "n_estimators": 200,
        "max_depth": 7,
        "learning_rate": 0.05,
        "class_weight": "balanced",
        "num_threads": CPU            # у LGB параметр другое имя
    },
    "hgb": {
        "max_iter": 200,
        "max_depth": 7,
        "learning_rate": 0.05,
        "random_state": 42
    }
}
estimators = [
    ("rf",  RandomForestClassifier(**MODEL_PARAMS["rf"])),
    ("xgb", XGBClassifier(**MODEL_PARAMS["xgb"])),
    ("lgb", LGBMClassifier(**MODEL_PARAMS["lgb"])),
    ("hgb", HistGradientBoostingClassifier(**MODEL_PARAMS["hgb"])),
]

strategy = SKLearnVotingClassificationStrategy(estimators, {"voting": "soft", "n_jobs": -1})
classifier = SpatialClassifier(strategy, 1000, 5)
score = classifier.train(russia)

classifier = SpatialClassifier.default()
result = classifier.run(russia)