from typing import Iterable

from funcy import juxt

from image_prediction.classifier.classifier import Classifier
from image_prediction.classifier.image_classifier import ImageClassifier
from image_prediction.compositor.compositor import TransformerCompositor
from image_prediction.encoder.encoders.hash_encoder import HashEncoder
from image_prediction.estimator.adapter.adapter import EstimatorAdapter
from image_prediction.formatter.formatter import format_image_plus
from image_prediction.formatter.formatters.camel_case import Snake2CamelCaseKeyFormatter
from image_prediction.formatter.formatters.enum import EnumFormatter
from image_prediction.image_extractor.extractors.parsable import ParsablePDFImageExtractor
from image_prediction.label_mapper.mappers.probability import ProbabilityMapper
from image_prediction.model_loader.loader import ModelLoader
from image_prediction.model_loader.loaders.mlflow import MlflowConnector
from image_prediction.redai_adapter.mlflow import MlflowModelReader
from image_prediction.transformer.transformers.coordinate.pdfnet import PDFNetCoordinateTransformer
from image_prediction.transformer.transformers.response import ResponseTransformer
from pdf2img.extraction import extract_images_via_metadata


def get_mlflow_model_loader(mlruns_dir):
    model_loader = ModelLoader(MlflowConnector(MlflowModelReader(mlruns_dir)))
    return model_loader


def get_image_classifier(model_loader, model_identifier):
    model, classes = juxt(model_loader.load_model, model_loader.load_classes)(model_identifier)
    return ImageClassifier(Classifier(EstimatorAdapter(model), ProbabilityMapper(classes)))


def get_dispatched_extract(**kwargs):
    image_extractor = ParsablePDFImageExtractor(**kwargs)

    def extract(pdf: bytes, page_range: range = None, metadata_per_image: Iterable[dict] = None):
        if metadata_per_image:
            image_pluses = extract_images_via_metadata(pdf, metadata_per_image)
            yield from map(format_image_plus, image_pluses)
        else:
            yield from image_extractor.extract(pdf, page_range)

    return extract


def get_formatter():
    formatter = TransformerCompositor(
        PDFNetCoordinateTransformer(), EnumFormatter(), ResponseTransformer(), Snake2CamelCaseKeyFormatter()
    )
    return formatter


def get_encoder():
    return HashEncoder()