import typing as T
from pathlib import Path
import json
import warnings
import logging

import pandas as pd
import torch
import torch.nn.functional as F
from torch.optim import lr_scheduler as optim_lr_scheduler
from torch_geometric.data import Data
from pytorch_lightning import LightningModule
from torchvision.ops import box_iou
from torchvision import transforms
import torchmetrics

from . import model_utils
from .cultio import CultioNet, GeoRefinement
from .maskcrnn import BFasterRCNN
from .base_layers import Softmax
from ..losses import TanimotoDistLoss

warnings.filterwarnings("ignore")
logging.getLogger("lightning").addHandler(logging.NullHandler())
logging.getLogger("lightning").propagate = False
logging.getLogger("lightning").setLevel(logging.ERROR)

class MaskRCNNLitModel(LightningModule):
    def __init__(
        self,
        cultionet_model_file: Path,
        cultionet_num_features: int,
        cultionet_num_time_features: int,
        cultionet_filters: int,
        cultionet_num_classes: int,
        ckpt_name: str = "maskrcnn",
        model_name: str = "maskrcnn",
        learning_rate: float = 1e-3,
        weight_decay: float = 1e-5,
        resize_height: int = 201,
        resize_width: int = 201,
        min_image_size: int = 100,
        max_image_size: int = 500,
        trainable_backbone_layers: int = 3,
    ):
        """Lightning model.

        Args:
            num_features
            num_time_features
            filters
            learning_rate
            weight_decay
        """
        super(MaskRCNNLitModel, self).__init__()
        self.save_hyperparameters()

        self.ckpt_name = ckpt_name
        self.model_name = model_name
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.num_classes = 2
        self.resize_height = resize_height
        self.resize_width = resize_width

        self.cultionet_model = CultioLitModel(
            num_features=cultionet_num_features,
            num_time_features=cultionet_num_time_features,
            filters=cultionet_filters,
            num_classes=cultionet_num_classes,
        )
        self.cultionet_model.load_state_dict(
            state_dict=torch.load(cultionet_model_file)
        )
        self.cultionet_model.eval()
        self.cultionet_model.freeze()
        self.model = BFasterRCNN(
            in_channels=4,
            out_channels=256,
            num_classes=self.num_classes,
            sizes=(16, 32, 64, 128, 256),
            aspect_ratios=(0.5, 1.0, 3.0),
            trainable_backbone_layers=trainable_backbone_layers,
            min_image_size=min_image_size,
            max_image_size=max_image_size,
        )

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

    def mask_forward(
        self,
        distance_ori: torch.Tensor,
        distance: torch.Tensor,
        edge: torch.Tensor,
        crop_r: torch.Tensor,
        height: T.Union[None, int, torch.Tensor],
        width: T.Union[None, int, torch.Tensor],
        batch: T.Union[None, int, torch.Tensor],
        y: T.Union[None, torch.Tensor] = None,
    ):
        height = int(height) if batch is None else int(height[0])
        width = int(width) if batch is None else int(width[0])
        batch_size = 1 if batch is None else batch.unique().size(0)
        x = torch.cat(
            (
                distance_ori,
                distance,
                edge[:, 1][:, None],
                crop_r[:, 1][:, None],
            ),
            dim=1,
        )
        # in x = (H*W x C)
        # new x = (B x C x H x W)
        gc = model_utils.GraphToConv()
        x = gc(x, batch_size, height, width)
        resizer = transforms.Resize((self.resize_height, self.resize_width))
        x = [resizer(image) for image in x]
        targets = None
        if y is not None:
            targets = []
            for bidx in y["image_id"].unique():
                batch_dict = {}
                batch_slice = y["image_id"] == bidx
                for k in y.keys():
                    if k == "masks":
                        batch_dict[k] = resizer(y[k][batch_slice])
                    elif k == "boxes":
                        # [xmin, ymin, xmax, ymax]
                        batch_dict[k] = self.scale_boxes(
                            y[k][batch_slice], batch, [height]
                        )
                    else:
                        batch_dict[k] = y[k][batch_slice]
                    targets.append(batch_dict)
        outputs = self.model(x, targets)

        return outputs

    def scale_boxes(
        self,
        boxes: torch.Tensor,
        batch: torch.Tensor,
        height: T.Union[None, int, T.List[int], torch.Tensor],
    ):
        height = int(height) if batch is None else int(height[0])
        scale = self.resize_height / height

        return boxes * scale

    def forward(
        self,
        batch: Data,
        batch_idx: int = None,
        y: T.Optional[torch.Tensor] = None,
    ) -> T.Tuple[
        torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
    ]:
        """Performs a single model forward pass."""
        with torch.no_grad():
            distance_ori, distance, edge, __, crop_r = self.cultionet_model(
                batch
            )
        estimates = self.mask_forward(
            distance_ori,
            distance,
            edge,
            crop_r,
            height=batch.height,
            width=batch.width,
            batch=batch.batch,
            y=y,
        )

        return estimates

    def on_save_checkpoint(self, checkpoint):
        """Save the checkpoint."""
        ckpt_file = Path(self.logger.save_dir) / f"{self.ckpt_name}.ckpt"
        if ckpt_file.is_file():
            ckpt_file.unlink()
        torch.save(checkpoint, ckpt_file)

    def on_validation_epoch_end(self, *args, **kwargs):
        """Save the model on validation end."""
        model_file = Path(self.logger.save_dir) / f"{self.model_name}.pt"
        if model_file.is_file():
            model_file.unlink()
        torch.save(self.state_dict(), model_file)

    def calc_loss(
        self, batch: T.Union[Data, T.List], y: T.Optional[torch.Tensor] = None
    ):
        """Calculates the loss for each layer.

        Returns:
            Average loss
        """
        losses = self(batch, y=y)
        loss = sum(loss for loss in losses.values())

        return loss

    def training_step(self, batch: Data, batch_idx: int = None):
        """Executes one training step."""
        y = {
            "boxes": batch.boxes,
            "labels": batch.box_labels,
            "masks": batch.box_masks,
            "image_id": batch.image_id,
        }
        loss = self.calc_loss(batch, y=y)
        self.log("loss", loss, on_step=False, on_epoch=True, prog_bar=True)

        return loss

    def _shared_eval_step(self, batch: Data) -> dict:
        # Predictions
        instances = self(batch)
        # True boxes
        true_boxes = self.scale_boxes(batch.boxes, batch, batch.height)

        predict_iou_score = torch.tensor(0.0, device=self.device)
        iou_score = torch.tensor(0.0, device=self.device)
        box_score = torch.tensor(0.0, device=self.device)
        for bidx, batch_value in enumerate(batch.image_id.unique()):
            # This should be low (i.e., low overlap of predicted boxes)
            predict_iou_score += box_iou(
                instances[bidx]["boxes"], instances[bidx]["boxes"]
            ).mean()
            # This should be high (i.e., high overlap of predictions and true boxes)
            iou_score += box_iou(
                true_boxes[batch.image_id == batch_value],
                instances[bidx]["boxes"],
            ).mean()
            # This should be high (i.e., masks should be confident)
            box_score += instances[bidx]["scores"].mean()
        predict_iou_score /= batch.image_id.unique().size(0)
        iou_score /= batch.image_id.unique().size(0)
        box_score /= batch.image_id.unique().size(0)

        total_iou_score = (predict_iou_score + (1.0 - iou_score)) * 0.5
        box_score = 1.0 - box_score
        # Minimize intersection-over-union and maximum score
        total_score = (total_iou_score + box_score) * 0.5

        metrics = {
            "predict_iou_score": predict_iou_score,
            "iou_score": iou_score,
            "box_score": box_score,
            "mean_score": total_score,
        }

        return metrics

    def validation_step(self, batch: Data, batch_idx: int = None) -> dict:
        """Executes one valuation step."""
        eval_metrics = self._shared_eval_step(batch)

        metrics = {
            "val_loss": eval_metrics["mean_score"],
            "val_piou": eval_metrics["predict_iou_score"],
            "val_iou": eval_metrics["iou_score"],
            "val_box": eval_metrics["box_score"],
        }
        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=True)

        return metrics

    def test_step(self, batch: Data, batch_idx: int = None) -> dict:
        """Executes one test step."""
        eval_metrics = self._shared_eval_step(batch)

        metrics = {
            "test_loss": eval_metrics["mean_score"],
            "test_piou": eval_metrics["predict_iou_score"],
            "test_iou": eval_metrics["iou_score"],
            "test_box": eval_metrics["box_score"],
        }
        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=True)

        return metrics

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            list(self.model.parameters()),
            lr=self.learning_rate,
            weight_decay=self.weight_decay,
            eps=1e-4,
        )
        lr_scheduler = optim_lr_scheduler.ReduceLROnPlateau(
            optimizer, factor=0.1, patience=5
        )

        return {
            "optimizer": optimizer,
            "scheduler": lr_scheduler,
            "monitor": "val_loss",
        }

def scale_logits(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    return x / t

class RefineLitModel(LightningModule):
    def __init__(
        self,
        in_features: int,
        num_classes: int = 2,
        learning_rate: float = 1e-3,
        weight_decay: float = 0.01,
        eps: float = 1e-4,
        edge_class: int = 2,
        class_counts: T.Optional[torch.Tensor] = None,
        cultionet_ckpt: T.Optional[T.Union[Path, str]] = None,
    ):
        super(RefineLitModel, self).__init__()

        self.save_hyperparameters()

        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.eps = eps
        self.edge_class = edge_class
        self.class_counts = class_counts
        self.cultionet_ckpt = cultionet_ckpt

        self.cultionet_model = None
        self.geo_refine_model = GeoRefinement(
            in_features=in_features, out_channels=num_classes
        )

        self.configure_loss()

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

    def forward(
        self,
        predictions: T.Dict[str, torch.Tensor],
        batch: Data,
        batch_idx: int = None,
    ) -> T.Dict[str, torch.Tensor]:
        return self.geo_refine_model(predictions, data=batch)

    def set_true_labels(self, batch: Data) -> torch.Tensor:
        # in case of multi-class, `true_crop` = 1, 2, etc.
        true_crop = torch.where(
            (batch.y > 0) & (batch.y != self.edge_class), 1, 0
        ).long()

        return true_crop

    def calc_loss(
        self,
        batch: T.Union[Data, T.List],
        predictions: T.Dict[str, torch.Tensor],
    ):
        true_crop = self.set_true_labels(batch)
        # Predicted crop values are probabilities
        loss = self.crop_loss(predictions["crop"], true_crop)

        return loss

    def training_step(
        self, batch: Data, batch_idx: int = None, optimizer_idx: int = None
    ):
        """Executes one training step."""
        # Apply inference with the main cultionet model
        if (self.cultionet_ckpt is not None) and (
            self.cultionet_model is None
        ):
            self.cultionet_model = CultioLitModel.load_from_checkpoint(
                checkpoint_path=str(self.cultionet_ckpt)
            )
            self.cultionet_model.to(self.device)
            self.cultionet_model.eval()
            self.cultionet_model.freeze()
        with torch.no_grad():
            predictions = self.cultionet_model(batch)

        predictions = self(predictions, batch)
        loss = self.calc_loss(batch, predictions)

        metrics = {"loss": loss}
        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=True)

        return metrics

    def on_train_epoch_end(self, *args, **kwargs):
        """Save the scaling parameters on training end."""
        if self.logger.save_dir is not None:
            model_file = Path(self.logger.save_dir) / "refine.pt"
            if model_file.is_file():
                model_file.unlink()
            torch.save(self.geo_refine_model.state_dict(), model_file)

    def configure_loss(self):
        self.crop_loss = TanimotoDistLoss(scale_pos_weight=True)

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            list(self.geo_refine_model.parameters()),
            lr=self.learning_rate,
            weight_decay=self.weight_decay,
            eps=self.eps,
        )
        lr_scheduler = optim_lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=20, eta_min=1e-5, last_epoch=-1
        )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": lr_scheduler,
                "monitor": "loss",
                "interval": "epoch",
                "frequency": 1,
            },
        }

class CultioLitModel(LightningModule):
    def __init__(
        self,
        num_features: int = None,
        num_time_features: int = None,
        num_classes: int = 2,
        filters: int = 32,
        model_type: str = "ResUNet3Psi",
        activation_type: str = "SiLU",
        dilations: T.Union[int, T.Sequence[int]] = None,
        res_block_type: str = "resa",
        attention_weights: str = "spatial_channel",
        optimizer: str = "AdamW",
        learning_rate: float = 1e-3,
        lr_scheduler: str = "CosineAnnealingLR",
        steplr_step_size: int = 5,
        weight_decay: float = 0.01,
        eps: float = 1e-4,
        ckpt_name: str = "last",
        model_name: str = "cultionet",
        deep_sup_dist: bool = False,
        deep_sup_edge: bool = False,
        deep_sup_mask: bool = False,
        class_counts: T.Optional[torch.Tensor] = None,
        edge_class: T.Optional[int] = None,
        temperature_lit_model: T.Optional[GeoRefinement] = None,
        scale_pos_weight: T.Optional[bool] = True,
        save_batch_val_metrics: T.Optional[bool] = False,
    ):
        """Lightning model."""
        super(CultioLitModel, self).__init__()

        self.save_hyperparameters()

        self.optimizer = optimizer
        self.learning_rate = learning_rate
        self.lr_scheduler = lr_scheduler
        self.steplr_step_size = steplr_step_size
        self.weight_decay = weight_decay
        self.eps = eps
        self.ckpt_name = ckpt_name
        self.model_name = model_name
        self.num_classes = num_classes
        self.num_time_features = num_time_features
        self.class_counts = class_counts
        self.temperature_lit_model = temperature_lit_model
        self.scale_pos_weight = scale_pos_weight
        self.save_batch_val_metrics = save_batch_val_metrics
        self.deep_sup_dist = deep_sup_dist
        self.deep_sup_edge = deep_sup_edge
        self.deep_sup_mask = deep_sup_mask
        self.sigmoid = torch.nn.Sigmoid()
        if edge_class is not None:
            self.edge_class = edge_class
        else:
            self.edge_class = num_classes

        self.model_attr = f"{model_name}_{model_type}"
        setattr(
            self,
            self.model_attr,
            CultioNet(
                ds_features=num_features,
                ds_time_features=num_time_features,
                filters=filters,
                num_classes=self.num_classes,
                model_type=model_type,
                activation_type=activation_type,
                dilations=dilations,
                res_block_type=res_block_type,
                attention_weights=attention_weights,
                deep_sup_dist=deep_sup_dist,
                deep_sup_edge=deep_sup_edge,
                deep_sup_mask=deep_sup_mask,
            ),
        )
        self.configure_loss()
        self.configure_scorer()

    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

    @property
    def cultionet_model(self) -> CultioNet:
        return getattr(self, self.model_attr)

    def forward(
        self, batch: Data, batch_idx: int = None
    ) -> T.Dict[str, torch.Tensor]:
        """Performs a single model forward pass.

        Returns:
            distance: Normalized distance transform (from boundaries), [0,1].
            edge: Probabilities of edge|non-edge, [0,1].
            crop: Logits of crop|non-crop.
        """
        return self.cultionet_model(batch)

    @staticmethod
    def get_cuda_memory():
        t = torch.cuda.get_device_properties(0).total_memory
        r = torch.cuda.memory_reserved(0)
        a = torch.cuda.memory_allocated(0)
        print(f"{t * 1e-6:.02f}MB", f"{r * 1e-6:.02f}MB", f"{a * 1e-6:.02f}MB")

    def predict_step(
        self, batch: Data, batch_idx: int = None
    ) -> T.Dict[str, torch.Tensor]:
        """A prediction step for Lightning."""
        predictions = self.forward(batch, batch_idx)
        if self.temperature_lit_model is not None:
            predictions = self.temperature_lit_model(predictions, batch)

        return predictions

    def get_true_labels(
        self, batch: Data, crop_type: torch.Tensor = None
    ) -> T.Dict[str, T.Union[None, torch.Tensor]]:
        """Gets true labels from the data batch."""
        true_edge = torch.where(batch.y == self.edge_class, 1, 0).long()
        # Recode all crop classes to 1, otherwise 0
        true_crop = torch.where(
            (batch.y > 0) & (batch.y < self.edge_class), 1, 0
        ).long()
        # Same as above, with additional edge class
        # Non-crop = 0
        # Crop | Edge = 1
        true_crop_and_edge = torch.where(batch.y > 0, 1, 0).long()
        # Same as above, with additional edge class
        # Non-crop = 0
        # Crop = 1
        # Edge = 2
        true_crop_or_edge = torch.where(
            (batch.y > 0) & (batch.y < self.edge_class),
            1,
            torch.where(batch.y == self.edge_class, 2, 0),
        ).long()
        true_crop_type = None
        if crop_type is not None:
            # Leave all crop classes as they are
            true_crop_type = torch.where(
                batch.y == self.edge_class, 0, batch.y
            ).long()

        return {
            "true_edge": true_edge,
            "true_crop": true_crop,
            "true_crop_and_edge": true_crop_and_edge,
            "true_crop_or_edge": true_crop_or_edge,
            "true_crop_type": true_crop_type,
        }

    def softmax(self, x: torch.Tensor, dim: int = 1) -> torch.Tensor:
        return F.softmax(x, dim=dim, dtype=x.dtype)

    def probas_to_labels(
        self, x: torch.Tensor, thresh: float = 0.5
    ) -> torch.Tensor:
        if x.shape[1] == 1:
            labels = x.gt(thresh).long()
        else:
            labels = x.argmax(dim=1).long()

        return labels

    def logits_to_probas(self, x: torch.Tensor) -> T.Union[None, torch.Tensor]:
        if x is not None:
            # Single-dimension inputs are sigmoid probabilities
            if x.shape[1] > 1:
                # Transform logits to probabilities
                x = self.softmax(x)
            else:
                x = self.sigmoid(x)
            x = x.clip(0, 1)

        return x

    # def on_train_epoch_start(self):
    #     # Get the current learning rate from the optimizer
    #     eps = self.optimizers().optimizer.param_groups[0]['eps']
    #     weight_decay = self.optimizers().optimizer.param_groups[0]['weight_decay']
    #     if (weight_decay != self.weight_decay) or (eps != self.eps):
    #         self.configure_optimizers()

    def on_validation_epoch_end(self, *args, **kwargs):
        """Save the model on validation end."""
        if self.logger.save_dir is not None:
            model_file = Path(self.logger.save_dir) / f"{self.model_name}.pt"
            if model_file.is_file():
                model_file.unlink()
            torch.save(self.state_dict(), model_file)

    def calc_loss(
        self,
        batch: T.Union[Data, T.List],
        predictions: T.Dict[str, torch.Tensor],
    ):
        """Calculates the loss.

        Returns:
            Total loss
        """
        true_labels_dict = self.get_true_labels(
            batch, crop_type=predictions["crop_type"]
        )

        # RNN level 2 loss (non-crop=0; crop|edge=1)
        crop_star_l2_loss = self.crop_star_l2_loss(
            predictions["crop_star_l2"], true_labels_dict["true_crop_and_edge"]
        )
        # RNN final loss (non-crop=0; crop=1; edge=2)
        crop_star_loss = self.crop_star_loss(
            predictions["crop_star"], true_labels_dict["true_crop_or_edge"]
        )
        # Main loss
        loss = (
            # RNN losses
            0.25 * crop_star_l2_loss
            + 0.5 * crop_star_loss
        )
        # Edge losses
        if self.deep_sup_dist:
            dist_loss_3_1 = self.dist_loss_3_1(
                predictions["dist_3_1"], batch.bdist
            )
            dist_loss_2_2 = self.dist_loss_2_2(
                predictions["dist_2_2"], batch.bdist
            )
            dist_loss_1_3 = self.dist_loss_1_3(
                predictions["dist_1_3"], batch.bdist
            )
            # Main loss
            loss = (
                loss
                + 0.1 * dist_loss_3_1
                + 0.25 * dist_loss_2_2
                + 0.5 * dist_loss_1_3
            )
        # Distance transform loss
        dist_loss = self.dist_loss(predictions["dist"], batch.bdist)
        # Main loss
        loss = loss + dist_loss
        # Distance transform losses
        if self.deep_sup_edge:
            edge_loss_3_1 = self.edge_loss_3_1(
                predictions["edge_3_1"], true_labels_dict["true_edge"]
            )
            edge_loss_2_2 = self.edge_loss_2_2(
                predictions["edge_2_2"], true_labels_dict["true_edge"]
            )
            edge_loss_1_3 = self.edge_loss_1_3(
                predictions["edge_1_3"], true_labels_dict["true_edge"]
            )
            # Main loss
            loss = (
                loss
                + 0.1 * edge_loss_3_1
                + 0.25 * edge_loss_2_2
                + 0.5 * edge_loss_1_3
            )
        # Edge loss
        edge_loss = self.edge_loss(
            predictions["edge"], true_labels_dict["true_edge"]
        )
        # Main loss
        loss = loss + edge_loss
        # Crop mask losses
        if self.deep_sup_mask:
            crop_loss_3_1 = self.crop_loss_3_1(
                predictions["crop_3_1"], true_labels_dict["true_crop"]
            )
            crop_loss_2_2 = self.crop_loss_2_2(
                predictions["crop_2_2"], true_labels_dict["true_crop"]
            )
            crop_loss_1_3 = self.crop_loss_1_3(
                predictions["crop_1_3"], true_labels_dict["true_crop"]
            )
            # Main loss
            loss = (
                loss
                + 0.1 * crop_loss_3_1
                + 0.25 * crop_loss_2_2
                + 0.5 * crop_loss_1_3
            )
        # Crop mask loss
        crop_loss = self.crop_loss(
            predictions["crop"], true_labels_dict["true_crop"]
        )
        # Main loss
        loss = loss + crop_loss

        if predictions["crop_type"] is not None:
            # Upstream (deep) loss on crop-type
            crop_type_star_loss = self.crop_type_star_loss(
                predictions["crop_type_star"],
                true_labels_dict["true_crop_type"],
            )
            loss = loss + crop_type_star_loss
            # Loss on crop-type
            crop_type_loss = self.crop_type_loss(
                predictions["crop_type"], true_labels_dict["true_crop_type"]
            )
            loss = loss + crop_type_loss

        return loss

    def training_step(self, batch: Data, batch_idx: int = None):
        """Executes one training step and logs training step metrics."""
        predictions = self(batch)
        loss = self.calc_loss(batch, predictions)
        self.log("loss", loss, on_step=False, on_epoch=True, prog_bar=True)

        return loss

    def _shared_eval_step(self, batch: Data, batch_idx: int = None) -> dict:
        predictions = self(batch)
        loss = self.calc_loss(batch, predictions)

        dist_mae = self.dist_mae(
            predictions["dist"].contiguous().view(-1),
            batch.bdist.contiguous().view(-1),
        )
        dist_mse = self.dist_mse(
            predictions["dist"].contiguous().view(-1),
            batch.bdist.contiguous().view(-1),
        )
        # Get the class labels
        edge_ypred = self.probas_to_labels(predictions["edge"])
        crop_ypred = self.probas_to_labels(predictions["crop"])
        # Get the true edge and crop labels
        true_labels_dict = self.get_true_labels(
            batch, crop_type=predictions["crop_type"]
        )
        # F1-score
        edge_score = self.edge_f1(edge_ypred, true_labels_dict["true_edge"])
        crop_score = self.crop_f1(crop_ypred, true_labels_dict["true_crop"])
        # MCC
        edge_mcc = self.edge_mcc(edge_ypred, true_labels_dict["true_edge"])
        crop_mcc = self.crop_mcc(crop_ypred, true_labels_dict["true_crop"])
        # Dice
        edge_dice = self.edge_dice(edge_ypred, true_labels_dict["true_edge"])
        crop_dice = self.crop_dice(crop_ypred, true_labels_dict["true_crop"])
        # Jaccard/IoU
        edge_jaccard = self.edge_jaccard(
            edge_ypred, true_labels_dict["true_edge"]
        )
        crop_jaccard = self.crop_jaccard(
            crop_ypred, true_labels_dict["true_crop"]
        )

        total_score = (
            loss
            + (1.0 - edge_score)
            + (1.0 - crop_score)
            + dist_mae
            + (1.0 - edge_mcc)
            + (1.0 - crop_mcc)
        )

        metrics = {
            "loss": loss,
            "dist_mae": dist_mae,
            "dist_mse": dist_mse,
            "edge_f1": edge_score,
            "crop_f1": crop_score,
            "edge_mcc": edge_mcc,
            "crop_mcc": crop_mcc,
            "edge_dice": edge_dice,
            "crop_dice": crop_dice,
            "edge_jaccard": edge_jaccard,
            "crop_jaccard": crop_jaccard,
            "score": total_score,
        }
        if predictions["crop_type"] is not None:
            crop_type_ypred = self.probas_to_labels(
                self.logits_to_probas(predictions["crop_type"])
            )
            crop_type_score = self.crop_type_f1(
                crop_type_ypred, true_labels_dict["true_crop_type"]
            )
            metrics["crop_type_f1"] = crop_type_score

        return metrics

    def validation_step(self, batch: Data, batch_idx: int = None) -> dict:
        """Executes one valuation step."""
        eval_metrics = self._shared_eval_step(batch, batch_idx)

        metrics = {
            "val_loss": eval_metrics["loss"],
            "vef1": eval_metrics["edge_f1"],
            "vcf1": eval_metrics["crop_f1"],
            "vmae": eval_metrics["dist_mae"],
            "val_score": eval_metrics["score"],
        }
        if "crop_type_f1" in eval_metrics:
            metrics["vctf1"] = eval_metrics["crop_type_f1"]

        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=True)

        if self.save_batch_val_metrics:
            self._save_batch_metrics(metrics, self.current_epoch, batch)

        return metrics

    def _save_batch_metrics(
        self, metrics: T.Dict[str, torch.Tensor], epoch: int, batch: Data
    ) -> None:
        """Saves batch metrics to a parquet file."""
        if not self.trainer.sanity_checking:
            write_metrics = {
                "epoch": [epoch] * len(batch.train_id),
                "train_ids": batch.train_id,
            }
            for k, v in metrics.items():
                write_metrics[k] = [float(v)] * len(batch.train_id)
            if self.logger.save_dir is not None:
                metrics_file = (
                    Path(self.logger.save_dir) / "batch_metrics.parquet"
                )
                if not metrics_file.is_file():
                    df = pd.DataFrame(write_metrics)
                    df.to_parquet(metrics_file)
                else:
                    df_new = pd.DataFrame(write_metrics)
                    df = pd.read_parquet(metrics_file)
                    df = pd.concat((df, df_new), axis=0)
                    df.to_parquet(metrics_file)

    def test_step(self, batch: Data, batch_idx: int = None) -> dict:
        """Executes one test step."""
        eval_metrics = self._shared_eval_step(batch, batch_idx)

        metrics = {
            "test_loss": eval_metrics["loss"],
            "tmae": eval_metrics["dist_mae"],
            "tmse": eval_metrics["dist_mse"],
            "tef1": eval_metrics["edge_f1"],
            "tcf1": eval_metrics["crop_f1"],
            "temcc": eval_metrics["edge_mcc"],
            "tcmcc": eval_metrics["crop_mcc"],
            "tedice": eval_metrics["edge_dice"],
            "tcdice": eval_metrics["crop_dice"],
            "tejaccard": eval_metrics["edge_jaccard"],
            "tcjaccard": eval_metrics["crop_jaccard"],
            "test_score": eval_metrics["score"],
        }
        if "crop_type_f1" in eval_metrics:
            metrics["tctf1"] = eval_metrics["crop_type_f1"]

        self.log_dict(metrics, on_step=False, on_epoch=True, prog_bar=True)

        return metrics

    def configure_scorer(self):
        self.dist_mae = torchmetrics.MeanAbsoluteError()
        self.dist_mse = torchmetrics.MeanSquaredError()
        self.edge_f1 = torchmetrics.F1Score(num_classes=2, average="micro")
        self.crop_f1 = torchmetrics.F1Score(num_classes=2, average="micro")
        self.edge_mcc = torchmetrics.MatthewsCorrCoef(num_classes=2)
        self.crop_mcc = torchmetrics.MatthewsCorrCoef(num_classes=2)
        self.edge_dice = torchmetrics.Dice(num_classes=2, average="micro")
        self.crop_dice = torchmetrics.Dice(num_classes=2, average="micro")
        self.edge_jaccard = torchmetrics.JaccardIndex(
            average="micro", num_classes=2
        )
        self.crop_jaccard = torchmetrics.JaccardIndex(
            average="micro", num_classes=2
        )
        if self.num_classes > 2:
            self.crop_type_f1 = torchmetrics.F1Score(
                num_classes=self.num_classes,
                average="weighted",
                ignore_index=0,
            )

    def configure_loss(self):
        self.dist_loss = TanimotoDistLoss()
        if self.deep_sup_dist:
            self.dist_loss_3_1 = TanimotoDistLoss()
            self.dist_loss_2_2 = TanimotoDistLoss()
            self.dist_loss_1_3 = TanimotoDistLoss()
        # Edge losses
        self.edge_loss = TanimotoDistLoss()
        if self.deep_sup_edge:
            self.edge_loss_3_1 = TanimotoDistLoss()
            self.edge_loss_2_2 = TanimotoDistLoss()
            self.edge_loss_1_3 = TanimotoDistLoss()
        # Crop mask losses
        self.crop_loss = TanimotoDistLoss(
            scale_pos_weight=self.scale_pos_weight
        )
        if self.deep_sup_mask:
            self.crop_loss_3_1 = TanimotoDistLoss(
                scale_pos_weight=self.scale_pos_weight
            )
            self.crop_loss_2_2 = TanimotoDistLoss(
                scale_pos_weight=self.scale_pos_weight
            )
            self.crop_loss_1_3 = TanimotoDistLoss(
                scale_pos_weight=self.scale_pos_weight
            )
        # Crop RNN losses
        self.crop_star_l2_loss = TanimotoDistLoss()
        self.crop_star_loss = TanimotoDistLoss()
        # FIXME:
        if self.num_classes > 2:
            self.crop_type_star_loss = TanimotoDistLoss(
                scale_pos_weight=self.scale_pos_weight
            )
            self.crop_type_loss = TanimotoDistLoss(
                scale_pos_weight=self.scale_pos_weight
            )

    def configure_optimizers(self):
        params_list = list(self.cultionet_model.parameters())
        if self.optimizer == "AdamW":
            optimizer = torch.optim.AdamW(
                params_list,
                lr=self.learning_rate,
                weight_decay=self.weight_decay,
                eps=self.eps,
            )
        elif self.optimizer == "SGD":
            optimizer = torch.optim.SGD(
                params_list,
                lr=self.learning_rate,
                weight_decay=self.weight_decay,
                momentum=0.9,
            )
        else:
            raise NameError("Choose either 'AdamW' or 'SGD'.")

        if self.lr_scheduler == "ExponentialLR":
            model_lr_scheduler = optim_lr_scheduler.ExponentialLR(
                optimizer, gamma=0.5
            )
        elif self.lr_scheduler == "CosineAnnealingLR":
            model_lr_scheduler = optim_lr_scheduler.CosineAnnealingLR(
                optimizer, T_max=20, eta_min=1e-5, last_epoch=-1
            )
        elif self.lr_scheduler == "StepLR":
            model_lr_scheduler = optim_lr_scheduler.StepLR(
                optimizer, step_size=self.steplr_step_size, gamma=0.5
            )
        else:
            raise NameError(
                "The learning rate scheduler is not implemented in Cultionet."
            )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": model_lr_scheduler,
                "name": "lr_sch",
                "monitor": "val_score",
                "interval": "epoch",
                "frequency": 1,
            },
        }


# -*- coding: utf-8 -*-
"""
Schema Extension Tests
"""

from pytest import mark, raises

from fudgeo.enumeration import ConstraintType, SQLFieldType
from fudgeo.extension.schema import (
    EnumerationConstraint, GlobConstraint, RangeConstraint, Schema)
from fudgeo.geopkg import Field, GeoPackage

@mark.parametrize('on_create, post_create', [
    (True, False),
    (False, False),
    (False, True),
    (True, True),
])
def test_create_geopackage(tmp_path, on_create, post_create):
    """
    Test create geopackage
    """
    path = tmp_path / 'test.gpkg'
    pkg = GeoPackage.create(path, enable_schema=on_create)
    fields = [Field(name='some_field_name', data_type=SQLFieldType.double),
              Field(name='bobloblaw', data_type=SQLFieldType.blob)]
    tbl = pkg.create_table(name='the_table', fields=fields)
    assert pkg.is_schema_enabled is on_create
    if not on_create:
        assert pkg.metadata is None
    else:
        assert isinstance(pkg.schema, Schema)
    if post_create:
        assert pkg.enable_schema_extension() is True
        assert pkg.is_schema_enabled is True
        assert isinstance(pkg.schema, Schema)
    tbl.drop()
    pkg.connection.close()
    if path.exists():
        path.unlink()
# End test_create_geopackage function

def test_add_constraints(tmp_path):
    """
    Add Constraints
    """
    path = tmp_path / 'test.gpkg'
    pkg = GeoPackage.create(path, enable_schema=True)
    constraint = RangeConstraint(name='sampleRange', min_value=1, max_value=10)
    constraints = [
        EnumerationConstraint(name='sampleEnum', values=[1, 3, 5, 7, 9]),
        GlobConstraint(name='sampleGlob', pattern='[1-2][0-9][0-9][0-9]'),
    ]
    pkg.schema.add_constraints(constraint)
    pkg.schema.add_constraints(constraints)
    cursor = pkg.connection.execute(
        """SELECT COUNT(1) AS C FROM gpkg_data_column_constraints""")
    count, = cursor.fetchone()
    assert count == 7
# End test_add_constraints function

@mark.parametrize('table_name, column_name, name, title, description, mime_type, constraint_name, msg', [
    ('table_name', 'column_name', None, None, None, None, None, 'table name "table_name" not found'),
    ('the_table', 'column_name', None, None, None, None, None, 'column name "column_name" not found'),
    ('the_table', 'bobloblaw', None, None, None, None, None, 'expected mime_type value for blob column'),
    ('the_table', 'some_field_name', None, None, None, None, 'SELECT', 'constraint name "SELECT" not found'),
])
def test_add_column_definition_validation(tmp_path, table_name, column_name, name,
                                          title, description, mime_type,
                                          constraint_name, msg):
    """
    Test add column definition validation steps
    """
    path = tmp_path / 'test.gpkg'
    pkg = GeoPackage.create(path, enable_schema=True)
    fields = [Field(name='some_field_name', data_type=SQLFieldType.double),
              Field(name='bobloblaw', data_type=SQLFieldType.blob)]
    pkg.create_table(name='the_table', fields=fields)
    schema = pkg.schema
    if msg:
        with raises(ValueError) as context:
            schema.add_column_definition(
                table_name, column_name, name, title, description,
                mime_type, constraint_name)
        assert context.value.args[0].startswith(msg)
    if constraint_name:
        constraint = GlobConstraint(name=constraint_name, pattern='[0-9]')
        schema.add_constraints(constraint)
        cursor = pkg.connection.execute("""SELECT * FROM gpkg_data_column_constraints""")
        records = cursor.fetchall()
        assert len(records) == 1
        schema.add_column_definition(
            table_name, column_name, name, title, description,
            mime_type, constraint_name)
# End test_add_column_definition_validation function

@mark.parametrize('name, values, description, count, exception', [
    ('a', [1, 2, 3], 'asdf', 3, None),
    ('b', (1, 2, 3), 'asdf', 3, None),
    ('c', (1, 2, 3, 3), 'asdf', 3, None),
    ('d', [], 'asdf', -1, ValueError),
    ('e', (), 'asdf', -1, ValueError),
    (None, (), 'asdf', -1, TypeError),
])
def test_enumeration_constraint(name, values, description, count, exception):
    """
    Test enumeration constraint
    """
    constraint = EnumerationConstraint(name, values, description)
    if exception:
        with raises(exception):
            constraint.validate()
    else:
        assert len(constraint.as_records()) == count
# End test_enumeration_constraint function

@mark.parametrize('name, pattern, description, exception', [
    ('a', '[0-9]', 'asdf', None),
    ('b', None, 'asdf', TypeError),
    ('c', 123, 'asdf', TypeError),
    ('d', Ellipsis, 'asdf', TypeError),
])
def test_glob_constraint(name, pattern, description, exception):
    """
    Test glob constraint
    """
    constraint = GlobConstraint(name, pattern, description)
    if exception:
        with raises(exception):
            constraint.validate()
    else:
        expected = [
            (ConstraintType.glob, name, pattern, None, None,
             None, None, description)]
        assert constraint.as_records() == expected
# End test_glob_constraint function

@mark.parametrize('name, min_value, max_value, min_inclusive, max_inclusive, description, exception', [
    ('a', 10, 100, True, True, 'asdf', None),
    ('b', 100, 10, True, True, 'asdf', ValueError),
    ('c', 100, 100, True, True, 'asdf', ValueError),
    ('d', None, 100, True, True, 'asdf', TypeError),
    ('e', 100, None, True, True, 'asdf', TypeError),
])
def test_range_constraint(name, min_value, max_value, min_inclusive, max_inclusive, description, exception):
    """
    Test range constraint
    """
    constraint = RangeConstraint(
        name, min_value, max_value, min_inclusive, max_inclusive, description)
    if exception:
        with raises(exception):
            constraint.validate()
    else:
        expected = [(ConstraintType.range_, name, None,
                     min_value, int(min_inclusive),
                     max_value, int(max_inclusive), description)]
        assert constraint.as_records() == expected
# End test_range_constraint function

if __name__ == '__main__':  # pragma: no cover
    pass


import torch
import torch.nn as nn
import torch.nn.functional as F
from models.layers.init_weights import init_weights

#===============================================================================
# SCM
#===============================================================================
class SCM(nn.Module):
    def __init__(self, in_channels, inter_channels):
        super(SCM, self).__init__()
        self.inter_channels = inter_channels
        if self.inter_channels is None:
            self.inter_channels = in_channels // 2
            if self.inter_channels == 0:
                self.inter_channels = 1

        self.conv_v_1 = nn.Conv2d(in_channels, self.inter_channels, kernel_size=1)
        self.conv_k_1 = nn.Conv2d(in_channels, self.inter_channels, kernel_size=1)
        self.conv_q_1 = nn.Conv2d(in_channels, self.inter_channels, kernel_size=1)
        self.conv_v_2 = nn.Conv2d(in_channels, self.inter_channels, kernel_size=1)
        self.conv_k_2 = nn.Conv2d(in_channels, self.inter_channels, kernel_size=1)
        self.conv_q_2 = nn.Conv2d(in_channels, self.inter_channels, kernel_size=1)

    def forward(self, x_1, x_2):
        V_1 = self.conv_v_1(x_1)
        K_1 = self.conv_k_1(x_1)
        Q_1 = self.conv_q_1(x_1)
        shape = V_1.shape
        # print(shape)
        V_1 = V_1.view(V_1.shape[0], V_1.shape[1], -1)
        K_1 = K_1.view(K_1.shape[0], K_1.shape[1], -1)
        Q_1 = Q_1.view(Q_1.shape[0], Q_1.shape[1], -1)

        V_2 = self.conv_v_2(x_2)
        K_2 = self.conv_k_2(x_2)
        Q_2 = self.conv_q_2(x_2)
        V_2 = V_2.view(V_2.shape[0], V_2.shape[1], -1)
        K_2 = K_2.view(K_2.shape[0], K_2.shape[1], -1)
        Q_2 = Q_2.view(Q_2.shape[0], Q_2.shape[1], -1)

        R12 = torch.sigmoid(torch.matmul(torch.transpose(Q_1, -1, -2), K_2))
        R21 = torch.sigmoid(torch.matmul(torch.transpose(Q_2, -1, -2), K_1))
        # print(R12.shape)

        Fu_1 = torch.matmul(R12, torch.transpose(V_2, -1, -2))
        Fu_2 = torch.matmul(R21, torch.transpose(V_1, -1, -2))
        Fu_1 = torch.transpose(Fu_1, -1, -2).view(shape)
        Fu_2 = torch.transpose(Fu_2, -1, -2).view(shape)
        # print(Fu_1.shape)

        y_1 = torch.cat([x_1, Fu_1], dim=1)
        y_2 = torch.cat([x_2, Fu_2], dim=1)
        # print(y_1.shape)
        return y_1, y_2

#===============================================================================
# PSP module
#===============================================================================
def _PSP1x1Conv(in_channels, out_channels, norm_layer, norm_kwargs):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, 1, bias=False),
        norm_layer(out_channels, **({} if norm_kwargs is None else norm_kwargs)),
        nn.ReLU(True)
    )

class PyramidPooling(nn.Module):
    def __init__(self, in_channels, **kwargs):
        super(PyramidPooling, self).__init__()
        out_channels = int(in_channels / 4)
        self.avgpool1 = nn.AdaptiveAvgPool2d(1)
        self.avgpool2 = nn.AdaptiveAvgPool2d(2)
        self.avgpool3 = nn.AdaptiveAvgPool2d(3)
        self.avgpool4 = nn.AdaptiveAvgPool2d(6)
        self.conv1 = _PSP1x1Conv(in_channels, out_channels, **kwargs)
        self.conv2 = _PSP1x1Conv(in_channels, out_channels, **kwargs)
        self.conv3 = _PSP1x1Conv(in_channels, out_channels, **kwargs)
        self.conv4 = _PSP1x1Conv(in_channels, out_channels, **kwargs)

    def forward(self, x):
        size = x.size()[2:]
        feat1 = F.interpolate(self.conv1(self.avgpool1(x)), size, mode='bilinear', align_corners=True)
        feat2 = F.interpolate(self.conv2(self.avgpool2(x)), size, mode='bilinear', align_corners=True)
        feat3 = F.interpolate(self.conv3(self.avgpool3(x)), size, mode='bilinear', align_corners=True)
        feat4 = F.interpolate(self.conv4(self.avgpool4(x)), size, mode='bilinear', align_corners=True)
        return torch.cat([x, feat1, feat2, feat3, feat4], dim=1)

#===============================================================================
# EADAM
#===============================================================================
class EADAM_V1(nn.Module):
    def __init__(self, in_channels, inter_channels=None):
        super().__init__()
        self.inter_channels = inter_channels
        if self.inter_channels is None:
            self.inter_channels = in_channels // 2
            if self.inter_channels == 0:
                self.inter_channels = 1

        self.conv_v = nn.Conv2d(in_channels, self.inter_channels, kernel_size=1)
        self.conv_k = nn.Conv2d(in_channels, self.inter_channels, kernel_size=1)
        self.conv_q = nn.Conv2d(in_channels, self.inter_channels, kernel_size=1)
        self.conv_out = nn.Conv2d(self.inter_channels, in_channels, kernel_size=1)

    def forward(self, x_en, x_de):
        V = self.conv_v(x_en)
        K = self.conv_k(x_en)
        Q = self.conv_q(x_de)
        shape = V.shape

        V = V.view(V.shape[0], V.shape[1], -1)
        K = K.view(V.shape[0], K.shape[1], -1)
        Q = Q.view(Q.shape[0], Q.shape[1], -1)

        H = torch.matmul(V, torch.transpose(K, -1, -2))
        H = torch.matmul(torch.sigmoid(H), Q) / (V.shape[1] * V.shape[2])**0.5

        H = H.view(shape)
        attmap = self.conv_out(H)

        return attmap + x_en + x_de

#===============================================================================
# multiscale input Module
#===============================================================================
class MultiInput(nn.Module):
    def __init__(self, in_channels, factor):
        super().__init__()
        self.conv_in = nn.Conv2d(3, in_channels, kernel_size=3, padding=1)
        self.bn = nn.BatchNorm2d(in_channels)
        self.conv_out = nn.Conv2d(in_channels * 2, in_channels, 1)
        self.factor = factor

    def forward(self, img, features):
        x = F.interpolate(img, scale_factor=self.factor, mode='bilinear', align_corners=True)
        x = self.conv_in(x)
        x = self.bn(x)
        out = torch.cat([x, features], dim=1)
        out = self.conv_out(out)

        return out

#===============================================================================
# TopHat Module
#===============================================================================
def dilation(x, size):
    maxpool = nn.MaxPool2d((size, size), stride=(1, 1), padding=size // 2)

    return maxpool(x)

def erosion(x, size):
    maxpool = nn.MaxPool2d((size, size), stride=(1, 1), padding=size // 2)

    return -maxpool(-x)

def black_tophat(x, size):
    x_dilate = dilation(x, size)
    x_close = erosion(x_dilate, size)

    return x_close - x

class TopHatBlockV2(nn.Module):
    def __init__(self, in_channels, poolsize=[3, 5, 7]):
        super().__init__()
        self.conv_in = nn.Conv2d(in_channels, in_channels, 1)
        self.conv_out = nn.Conv2d(in_channels, in_channels, 1)

        self.poolsize = poolsize
        self.conv_1x1 = nn.Conv2d(in_channels * len(poolsize), in_channels, 1)

    def forward(self, x):
        identity = x

        x = self.conv_in(x)

        out = []
        for idx, size in enumerate(self.poolsize):
            out.append(black_tophat(x, size))

        out = torch.cat(out, dim=1)
        out = self.conv_1x1(out)

        return out + identity

#===============================================================================
# PCA Blocks
#===============================================================================
class SAPCABlockV5(nn.Module):
    """
    Spatial attention block
    """

    def __init__(self, in_channels, scale=3):
        super().__init__()
        self.scale = scale
        self.conv1 = nn.Conv2d(in_channels, in_channels, 1)
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, 3, padding=1, bias=False),
            nn.BatchNorm2d(in_channels),
        )

    def forward(self, x):
        gx = self.conv1(x)
        gx_ = gx.view(gx.shape[0], gx.shape[1], -1)
        gx_ = gx_ - torch.mean(gx_, dim=-1, keepdim=True)

        cmat = torch.matmul(gx_, gx_.transpose(-1, -2)) / gx.shape[0]
        w = []
        for i in range(cmat.shape[0]):
            eigval, eigvec = torch.symeig(cmat[i], eigenvectors=True)
            w.append(eigvec[:, -16:].transpose(-1, -2))
        w = torch.stack(w, dim=0)

        z = torch.matmul(w, gx_)  # 这一步可以考虑加softmax
        z = F.softmax(z * self.scale, dim=1)
        y_ = torch.matmul(w.transpose(-1, -2), z)
        y_ = y_.view(gx.shape[0], gx.shape[1], gx.shape[2], gx.shape[3])

        attmap = self.conv2(y_)

        x_en = x + attmap

        return x_en

class SAPCABlockV4(nn.Module):
    """
    Spatial attention block
    """

    def __init__(self, in_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, in_channels, 1)
        self.conv2 = nn.Sequential(nn.Conv2d(in_channels, in_channels, 1, bias=False),
                                   # norm_layer(in_channels),
                                   )

    def forward(self, x):
        gx = self.conv1(x)
        gx_ = gx.view(gx.shape[0], gx.shape[1], -1)
        gx_ = gx_ - torch.mean(gx_, dim=-1, keepdim=True)

        cmat = torch.matmul(gx_, gx_.transpose(-1, -2)) / gx.shape[0]
        w = []
        for i in range(cmat.shape[0]):
            eigval, eigvec = torch.symeig(cmat[i], eigenvectors=True)
            w.append(eigvec[:, -x.shape[1] // 4:].transpose(-1, -2))
        w = torch.stack(w, dim=0)

        z = torch.matmul(w, gx_)  # 这一步可以考虑加softmax
        y_ = torch.matmul(w.transpose(-1, -2), z)
        y_ = y_.view(gx.shape[0], gx.shape[1], gx.shape[2], gx.shape[3])

        attmap = self.conv2(y_)

        x_en = x + attmap

        return x_en

#================================================================
#================================================================
#================================================================
class SAPCABlockV3(nn.Module):
    """
    # Spatial attention block
    """

    def __init__(self, in_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, in_channels, 1)
        self.conv2 = nn.Conv2d(in_channels // 4, in_channels, 1)
        self.prelu = nn.PReLU()

    def forward(self, x):
        gx = self.conv1(x)
        gx_ = gx.view(gx.shape[0], gx.shape[1], -1)
        gx_ = gx_ - torch.mean(gx_, dim=-1, keepdim=True)

        cmat = torch.matmul(gx_, gx_.transpose(-1, -2)) / gx.shape[0]
        w = []
        for i in range(cmat.shape[0]):
            eigval, eigvec = torch.symeig(cmat[i], eigenvectors=True)
            w.append(eigvec[:, -x.shape[1] // 4:].transpose(-1, -2))
        w = torch.stack(w, dim=0)

        y = torch.matmul(w, gx_)
        y = y.view(y.shape[0], y.shape[1], gx.shape[2], gx.shape[3])

        attmap = self.conv2(y)
        x_en = x + attmap

        return x_en

class SAPCABlockV2(nn.Module):
    """
    #Spatial attention block
    """

    def __init__(self, in_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, in_channels, 1)
        self.conv2 = nn.Conv2d(in_channels // 4, in_channels, 1)
        self.prelu = nn.PReLU()

    def forward(self, x):
        gx = self.prelu(self.conv1(x))
        gx_ = gx.view(gx.shape[0], gx.shape[1], -1)

        cmat = torch.matmul(gx_, gx_.transpose(-1, -2)) / gx.shape[0]
        w = []
        for i in range(cmat.shape[0]):
            eigval, eigvec = torch.symeig(cmat[i], eigenvectors=True)
            w.append(eigvec[:, -x.shape[1] // 4:].transpose(-1, -2))
        w = torch.stack(w, dim=0)

        y = torch.matmul(w, gx_)
        y = torch.sigmoid(y)
        y = y.view(y.shape[0], y.shape[1], gx.shape[2], gx.shape[3])

        attmap = self.conv2(y)
        x_en = x * torch.tanh(attmap)

        return x_en

class SAPCABlock(nn.Module):
    """
    #Spatial attention block
    """

    def __init__(self, in_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, in_channels, 1)
        self.conv2 = nn.Conv2d(in_channels // 4, in_channels, 1)
        self.prelu = nn.PReLU()

    def forward(self, x):
        gx = self.prelu(self.conv1(x))
        gx_ = gx.view(gx.shape[0], gx.shape[1], -1)

        cmat = torch.matmul(gx_, gx_.transpose(-1, -2)) / gx.shape[0]
        w = []
        for i in range(cmat.shape[0]):
            eigval, eigvec = torch.symeig(cmat[i], eigenvectors=True)
            w.append(eigvec[:, -x.shape[1] // 4:].transpose(-1, -2))
        w = torch.stack(w, dim=0)

        y = torch.matmul(w, gx_)
        y = y.view(y.shape[0], y.shape[1], gx.shape[2], gx.shape[3])

        attmap = self.conv2(y)
        x_en = x + attmap

        return x_en

class DecoderBlock(nn.Module):
    def __init__(self, in_channels, n_filters):
        super(DecoderBlock, self).__init__()

        self.conv1 = nn.Conv2d(in_channels, in_channels // 4, 1)
        self.norm1 = nn.BatchNorm2d(in_channels // 4)
        self.relu1 = nn.ReLU(inplace=True)

        self.deconv2 = nn.ConvTranspose2d(in_channels // 4, in_channels // 4, 3, stride=2, padding=1, output_padding=1)
        self.norm2 = nn.BatchNorm2d(in_channels // 4)
        self.relu2 = nn.ReLU(inplace=True)

        self.conv3 = nn.Conv2d(in_channels // 4, n_filters, 1)
        self.norm3 = nn.BatchNorm2d(n_filters)
        self.relu3 = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv1(x)
        x = self.norm1(x)
        x = self.relu1(x)
        x = self.deconv2(x)
        x = self.norm2(x)
        x = self.relu2(x)
        x = self.conv3(x)
        x = self.norm3(x)
        x = self.relu3(x)

        return x


from nltk.stem import WordNetLemmatizer
from nltk.corpus import stopwords
# uncomment these lines to download the NLTK packages
#nltk.download('punkt')
#nltk.download('wordnet')
#nltk.download('stopwords')
#nltk.download('omw-1.4')

import argparse
import gensim
import numpy as np
from os.path import exists
from typing import List, Dict
from collections import defaultdict
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt

from utils.data import load_data, get_word_tokenized_corpus, get_data_property, get_data_chunks, ABSTRACT, N_CITATION
from utils.embeddings import train_fasttext_embedding, pretrain_fasttext_embedding, get_chunk_embeddings, save_fasttext, load_fasttext
from utils.features import get_features
from utils.controls import get_controls

def setup_args() -> argparse.Namespace:
    """Gets the command-line arguments used. Feel free to add your own here.

    Returns:
        argparse.Namespace: argparse object
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_name", type=str, default='fasttext_model/wiki-news-300d-1M.vec', help="File to load model from.")
    parser.add_argument("--train_model", action='store_true', help="Whether to train existing model on data.")

    parser.add_argument("--data_file", type=str, default='data/dblp-ref-0.json', help="File/directory to load data from.")
    parser.add_argument("--data_file_type", type=str, help="If a directory is provided as the datafile, provide type to read")
    parser.add_argument("--chunk_embs_file", type=str, default='data/chunk_embs.txt', help="File to save/load chunks from.")
    parser.add_argument("--proj_dir", type=str, default='./word_embedding_measures/', help="Directory storing all data/models.")

    parser.add_argument("--limit", type=int, default=30000, help="Number of documents to train on")
    # TODO: allow user to create chunks based on length (e.g. I want a chunk to be a sentence or I want a chunk to be 5 tokens)
    parser.add_argument("--T", type=int, default=20, help="Number of chunks to make")
    args = parser.parse_args()

    if args.data_file[-1] == '/':
        assert args.data_file_type is not None
    return args

def setup(args: argparse.Namespace) -> tuple:
    """Sets up the FastText model (either training or just loading) and
    loads in and preprocesses the data.

    Args:
        args(argparse.Namespace): argparse object

    Returns:
        tuple[gensim.models.fasttext.FastText, List[Dict]]: returns FastText model and data
    """
    proj_dir = args.proj_dir
    stemmer = WordNetLemmatizer()
    en_stop = set(stopwords.words('english'))

    print('Loading data...')
    data = load_data(args)
    documents = get_data_property(data, ABSTRACT)

    ft_model = None
    if exists(proj_dir + args.model_name):
        print('Retrieving model...')
        ft_model = load_fasttext(proj_dir + args.model_name)

        if args.train_model:
            print('Training model...')
            tokenized_data = get_word_tokenized_corpus(documents, stemmer, en_stop)
            ft_model = train_fasttext_embedding(ft_model, tokenized_data)
            save_fasttext(ft_model, proj_dir + args.model_name.replace('.bin', '-trained.bin'))
    else:
        # tokenize data and train fasttext model
        print('Training model...')
        tokenized_data = get_word_tokenized_corpus(documents, stemmer, en_stop)
        ft_model = pretrain_fasttext_embedding(tokenized_data)
        save_fasttext(ft_model, proj_dir + args.model_name)

    return ft_model, data

def setup_chunk_embeddings(args: argparse.Namespace, ft_model: gensim.models.fasttext.FastText, documents: List[str]) -> np.ndarray:
    """_summary_

    Args:
        args (argparse.Namespace): argparse object
        ft_model (gensim.models.fasttext.FastText): fasttext model
        documents (List[str]): list of documents

    Returns:
        np.ndarray: returns numpy array of chunk embeddings
    """

    if exists(args.proj_dir + args.chunk_embs_file):
        print('Loading chunk embeddings...')
        with open(args.proj_dir + args.chunk_embs_file, 'r+') as f:
            shape = tuple(map(int, f.readline()[1:].split(',')))
            limit = min(shape[0], args.limit)
            chunk_embs = np.loadtxt(f, skiprows=0, delimiter=',', max_rows=limit * args.T).reshape(limit, shape[1], shape[2])
    else:
        # should be of size [# documents, T, len(document)/T]
        print('Chunking documents...')
        chunks = [get_data_chunks(document, T=args.T) for document in documents]
        chunk_embs = np.array([get_chunk_embeddings(ft_model, chunk) for chunk in chunks])

        header = ','.join(map(str, chunk_embs.shape))
        np.savetxt(args.proj_dir + args.chunk_embs_file, chunk_embs.reshape(-1, chunk_embs.shape[-1]), header=header, delimiter=',')

    return chunk_embs

if __name__ == "__main__":
    args = setup_args()

    # additional properties I want from the data
    strict_loading_list = [ABSTRACT, N_CITATION, 'year', 'venue']
    args.strict_loading_list = strict_loading_list

    if exists(args.proj_dir + args.chunk_embs_file):
        print('Found existing chunks, loading data...')
        data = load_data(args)
        chunk_embs = setup_chunk_embeddings(args, None, None)
    else:
        print('No existing chunks, calling setup...')
        ft_model, data = setup(args)
        documents = get_data_property(data, ABSTRACT)
        chunk_embs = setup_chunk_embeddings(args, ft_model, documents)

    # this is the dependent variable we will do regression on
    # you should change this to whatever metric you want
    citation_counts = get_data_property(data, N_CITATION)
    # controls are very important for the regression and should be changed
    # based on your own use case
    controls = get_controls(data)

    print('Getting features...')
    features = [get_features(chunk_emb) for chunk_emb in chunk_embs]
    features_dict = defaultdict(list)

    # move features into a dict that stores lists for every feature
    for d in features:
        for k, v in d.items():
            features_dict[k].append(v)

    '''
    for key, value in features_dict.items():
        print(key)
        print(value)
        print('-'*100)
    '''

    # if there is a missing value at the end of distances (sometimes there are only T-1 chunks)
    for i, l in enumerate(features_dict['distances']):
        for j in range(args.T - 1 - len(l)):
            l.append(np.nan)

    avg_distances = np.nanmean(np.array(features_dict['distances']), axis=0, dtype='float32')
    # plots the speed over the chunks
    plt.plot(list(range(args.T-1)), avg_distances)
    plt.savefig('distances.png')

    print('Getting coefficients...')
    for key, value in features_dict.items():
        if key == 'distances':
            continue

        # remove rows with invalid values
        nan_vals = np.argwhere(np.isnan(value))

        controls = np.delete(controls, nan_vals, axis=0)
        non_nan_citations = np.log(1 + np.delete(np.array(citation_counts), nan_vals))
        non_nan_vals = np.log(np.delete(value, nan_vals)).reshape(-1, 1)
        dependent_vars = np.concatenate((non_nan_vals, controls), axis=1)

        clf = LinearRegression()
        clf.fit(dependent_vars, non_nan_citations)
        print(f'{key} coeff {clf.coef_[0]} mean {np.mean(non_nan_vals)}')


"""
Function helpers to manage PDF files.
"""
from PyPDF2 import PdfFileMerger, PdfFileReader

from docstamp.commands import call_command

def merge_pdfs(pdf_filepaths, out_filepath):
    """ Merge all the PDF files in `pdf_filepaths` in a new PDF file `out_filepath`.

    Parameters
    ----------
    pdf_filepaths: list of str
        Paths to PDF files.

    out_filepath: str
        Path to the result PDF file.

    Returns
    -------
    path: str
        The output file path.
    """
    merger = PdfFileMerger()
    for pdf in pdf_filepaths:
        merger.append(PdfFileReader(open(pdf, 'rb')))

    merger.write(out_filepath)

    return out_filepath

def pdf_to_cmyk(input_file, output_file):
    """ User `gs` (Ghostscript) to convert the colore model of a PDF to CMYK.

    Parameters
    ----------
    input_file: str

    output_file: str
    """
    cmd_args = [
        '-dSAFER',
        '-dBATCH',
        '-dNOPAUSE',
        '-dNOCACHE',
        '-sDEVICE=pdfwrite',
        '-sColorConversionStrategy=CMYK',
        '-dProcessColorModel=/DeviceCMYK',
        '-sOutputFile="{}" "{}"'.format(output_file, input_file),
    ]
    call_command('gs', cmd_args)


"""Helper functions to work with SAM GraphQLApi resource
"""

from typing import Any, Dict, List, Tuple, Union

SCHEMA_ARTIFACT_PROPERTY = "SchemaUri"
CODE_ARTIFACT_PROPERTY = "CodeUri"

def find_all_paths_and_values(property_name: str, graphql_dict: Dict[str, Any]) -> List[Tuple[str, Union[str, Dict]]]:
    """Find paths to the all properties with property_name and their (properties) values.

    It leverages the knowledge of GraphQLApi structure instead of doing generic search in the graph.

    Parameters
    ----------
    property_name
        Name of the property to look up, for example 'CodeUri'
    graphql_dict
        GraphQLApi resource dict

    Returns
    -------
        list of tuple (path, value) for all found properties which has property_name
    """
    # need to look up only in "Resolvers" and "Functions" subtrees
    resolvers_and_functions = {k: graphql_dict[k] for k in ("Resolvers", "Functions") if k in graphql_dict}
    stack: List[Tuple[Dict[str, Any], str]] = [(resolvers_and_functions, "")]
    paths_values: List[Tuple[str, Union[str, Dict]]] = []

    while stack:
        node, path = stack.pop()
        if isinstance(node, dict):
            for key, value in node.items():
                if key == property_name:
                    paths_values.append((f"{path}{key}", value))
                elif isinstance(value, dict):
                    stack.append((value, f"{path}{key}."))
        # there is no need to handle lists because
        # paths to "CodeUri" within "Resolvers" and "Functions" doesn't have lists
    return paths_values


import torch
import torch.optim
from collections import OrderedDict
import torch.nn as nn

class Net(nn.Module):
    def __init__(self, seq_net, name='MLP', activation=torch.tanh):
        super().__init__()
        self.features = OrderedDict()
        for i in range(len(seq_net) - 1):
            self.features['{}_{}'.format(name, i)] = nn.Linear(seq_net[i], seq_net[i + 1], bias=True)
        self.features = nn.ModuleDict(self.features)
        self.active = activation

        # initial_bias

        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.constant_(m.bias, 0)
                nn.init.xavier_normal_(m.weight, gain=1.0) #ini2
                # torch.nn.init.uniform_(m.weight, a=-0.01, b=0.01) #3
                # torch.nn.init.uniform_(m.weight, a=-0.1, b=0.1) #4
                # torch.nn.init.normal_(m.weight, mean=0.0, std=0.01) #5
                # torch.nn.init.normal_(m.weight, mean=0.0, std=0.1) #6

    def forward(self, x):
        # x = x.view(-1, 2)
        length = len(self.features)
        i = 0
        for name, layer in self.features.items():
            x = layer(x)
            if i == length - 1: break
            i += 1
            x = self.active(x)
        return x


import logging
import os

import numpy as np

from gym import spaces
from .nes_env import NesEnv, MetaNesEnv

logger = logging.getLogger(__name__)

# (world_number, level_number, area_number, max_distance)
WORLD_NUMBER = 0
LEVEL_NUMBER = 1
AREA_NUMBER = 2
MAX_DISTANCE = 3
SMB_LEVELS = [
    (1, 1, 1, 3266), (1, 2, 2, 3266), (1, 3, 4, 2514), (1, 4, 5, 2430),
    (2, 1, 1, 3298), (2, 2, 2, 3266), (2, 3, 4, 3682), (2, 4, 5, 2430),
    (3, 1, 1, 3298), (3, 2, 2, 3442), (3, 3, 3, 2498), (3, 4, 4, 2430),
    (4, 1, 1, 3698), (4, 2, 2, 3266), (4, 3, 4, 2434), (4, 4, 5, 2942),
    (5, 1, 1, 3282), (5, 2, 2, 3298), (5, 3, 3, 2514), (5, 4, 4, 2429),
    (6, 1, 1, 3106), (6, 2, 2, 3554), (6, 3, 3, 2754), (6, 4, 4, 2429),
    (7, 1, 1, 2962), (7, 2, 2, 3266), (7, 3, 4, 3682), (7, 4, 5, 3453),
    (8, 1, 1, 6114), (8, 2, 2, 3554), (8, 3, 3, 3554), (8, 4, 4, 4989)]
SUPER_MARIO_ROM_PATH = os.path.join(os.path.dirname(__file__), 'roms', 'super-mario.nes')

# --------------
# Helper Methods
# --------------
def is_int16(str):
    try:
        int(str, 16)
        return True
    except ValueError:
        return False

# --------------
# Classes
# --------------
class SuperMarioBrosEnv(NesEnv):
    def __init__(self, draw_tiles=False, level=0):
        NesEnv.__init__(self)
        package_directory = os.path.dirname(os.path.abspath(__file__))
        self.level = level
        self.draw_tiles = 1 if draw_tiles else 0
        self._mode = 'algo'             # 'algo' or 'human'
        self.lua_path.append(os.path.join(package_directory, 'lua/super-mario-bros.lua'))
        self.tiles = None
        self.launch_vars['target'] = self._get_level_code(self.level)
        self.launch_vars['mode'] = 'algo'
        self.launch_vars['meta'] = '0'
        self.launch_vars['draw_tiles'] = str(self.draw_tiles)
        if os.path.isfile(SUPER_MARIO_ROM_PATH):
            self.rom_path = SUPER_MARIO_ROM_PATH

        self.additional_rewards = {
            "time" : -0.1, #per second that passes by
            "death" : -100., #mario dies
            "extra_life" : 100., #mario gets an extra life, which includes getting 100th coin
            "mushroom" : 20., #mario eats a mushroom to become big
            "flower" : 25., #mario eats a flower
            "mushroom_hit" : -10., #mario gets hit while big
            "flower_hit" : -15., #mario gets hit while fire mario
            "coin" : 1.0, #mario gets a coin
            "finish_level" : 100.,
            "score" : 0.1, #per 100 points
        }

        # Tile mode
        if 1 == self.draw_tiles:
            self.tile_height = 13
            self.tile_width = 16
            self.screen_height = 13
            self.screen_width = 16
            self.tiles = np.zeros(shape=(self.tile_height, self.tile_width), dtype=np.uint8)
            self.observation_space = spaces.Box(low=0, high=3, shape=(self.tile_height, self.tile_width))

    # --------------
    # Properties
    # --------------
    @property
    def mode(self):
        return self._mode

    @mode.setter
    def mode(self, value):
        self._mode = value
        self.launch_vars['mode'] = value
        self.cmd_args = ['--xscale 2', '--yscale 2', '-f 0']
        if 'human' == value:
            self.disable_out_pipe = True
            self.disable_in_pipe = True
        else:
            self.disable_out_pipe = False
            self.disable_in_pipe = False

    # --------------
    # Methods
    # --------------
    def _get_level_code(self, level):
        world_number = int(level / 4) + 1
        level_number = (level % 4) + 1
        area_number = level_number
        # Worlds 1, 2, 4, 7 have a transition as area number 2 (so 2-2 is area 3 and 3, 2-3 is area 4, 2-4 is area 5)
        if world_number in [1, 2, 4, 7] and level_number >= 2:
            area_number += 1
        return '%d%d%d' % (world_number, level_number, area_number)

    #This function has been heavily modified to make the reward involve more than just raw x distance
    #These data messages are defined in the lua file for super-mario-bros

    #TODO - need to edit LUA file to include star power rewards

    def _process_data_message(self, frame_number, data):
        # Format: data_<frame>#name_1:value_1|name_2:value_2|...
        if frame_number <= self.last_frame or self.info is None:
            return
        parts = data.split('|')

        #rewards have to do with x distance, coins, lifes, mushrooms and fireballs, and score. we construct them one by one
        self.reward = 0.
        if not self.episode_reward:
            self.episode_reward = 0.

        for part in parts:
            if part.find(':') == -1:
                continue
            parts_2 = part.split(':')
            name = parts_2[0]
            value = int(parts_2[1])

            """
            The following rewards are unimplemented, by choice
            self.additional_rewards = {
                "finish_level" : 100.
                "score" : 0.1 #per 100 points
            }
            """

            #detect if finished
            if 'is_finished' == name:
                self.is_finished = bool(value)

            #detect extra lives or deaths
            elif 'life' == name:
                #print("life change detected. life value is now %d" % value)
                self.info[name] = value
                if value < 3:
                    #print("detected death")
                    self.reward += self.additional_rewards['death']
                    #self.episode_reward += self.additional_rewards['death']
                elif value > 3:
                    #print("new life detected")
                    self.reward += self.additional_rewards['extra_life']
                    #self.episode_reward += self.additional_rewards['extra_life']

            #detect getting a coin
            elif 'coin' == name:
                #print("coin detected")
                self.reward += self.additional_rewards['coin']
                self.episode_reward += self.additional_rewards['coin']
                self.info[name] = value

            #detect if we decremented by a mario second, in which case, apply the reward
            elif 'time' == name and value < self.info[name]:
                #print("Time decrement detected")
                self.reward += self.additional_rewards['time']
                self.episode_reward += self.additional_rewards['time']
                self.info[name] = value

            #detect if distance incremented
            elif 'distance' == name:
                self.reward += value - self.info[name] #change in x distance
                self.episode_reward += value #total x distance
                self.info[name] = value #total x distance

            #detect if mario ate a mushroom, ate a flower, or got hit without dying
            elif 'player_status' == name:

                #2 - fire mario. only achieved if eating flower while super mario
                #1 - super mario. only achieved if eating mushroom while small mario
                #0 - small mario. only achieved if hit while super mario or fire mario. if hit while small mario, death.

                #ate a flower (assuming was still super mario. if eating flower while small mario, mario only becomes super mario so this value
                #would be a value of 1, and be caught in the value == 1 checks)
                if value == 2 and value > self.info[name]:
                    self.reward += self.additional_rewards['flower']
                    self.episode_reward += self.additional_rewards['flower']

                #if currently super mario, only need to check if this is from eating mushroom. if hit while fire mario, goes back to small mario
                elif value == 1 and value > self.info[name]:
                    self.reward += self.additional_rewards['mushroom']
                    self.episode_reward += self.additional_rewards['mushroom']

                #if small value was sent, you got hit when you were big
                elif value == 0 and self.info[name] == 1:
                    self.reward += self.additional_rewards['mushroom_hit']
                    self.episode_reward += self.additional_rewards['mushroom_hit']

                #or worse, you got hit when you were a flower
                elif value == 0 and self.info[name] == 2:
                    self.reward += self.additional_rewards['flower_hit']
                    self.episode_reward += self.additional_rewards['flower_hit']

                self.info[name] = value

            else:
                self.info[name] = value

    def _process_screen_message(self, frame_number, data):
        # Format: screen_<frame>#<x (2 hex)><y (2 hex)><palette (2 hex)>|<x><y><p>|...
        if frame_number <= self.last_frame or self.screen is None:
            return
        parts = data.split('|')
        for part in parts:
            if 6 == len(part) and is_int16(part[0:2]) and is_int16(part[2:4]):
                x = int(part[0:2], 16)
                y = int(part[2:4], 16)
                self.screen[y][x] = self._get_rgb_from_palette(part[4:6])

    def _process_tiles_message(self, frame_number, data):
        # Format: tiles_<frame>#<x (1 hex)><y (1 hex)><value (1 hex)>|<x><y><v>|...
        if frame_number <= self.last_frame or self.tiles is None:
            return
        parts = data.split('|')
        for part in parts:
            if 3 == len(part) and is_int16(part[0:1]) and is_int16(part[1:2]) and is_int16(part[2:3]):
                x = int(part[0:1], 16)
                y = int(part[1:2], 16)
                v = int(part[2:3], 16)
                self.tiles[y][x] = v
                if v == 0: self.screen[y][x] = self._get_rgb_from_palette('0D')
                if v == 1: self.screen[y][x] = self._get_rgb_from_palette('30')
                if v == 2: self.screen[y][x] = self._get_rgb_from_palette('27')
                if v == 3: self.screen[y][x] = self._get_rgb_from_palette('05')

    def _process_ready_message(self, frame_number):
        # Format: ready_<frame>
        if 0 == self.last_frame:
            self.last_frame = frame_number

    def _process_done_message(self, frame_number):
        # Done means frame is done processing, please send next command
        # Format: done_<frame>
        if frame_number > self.last_frame:
            self.last_frame = frame_number

    def _process_reset_message(self):
        # Reset means 'changelevel' needs to be sent and last_frame needs to be set to 0
        # Not implemented in non-meta levels
        pass

    def _process_exit_message(self):
        # Exit means fceux is terminating
        # Format: exit
        self.is_finished = True
        self._is_exiting = 1
        self.close()

    def _parse_frame_number(self, parts):
        # Parsing frame number
        try:
            frame_number = int(parts[1]) if len(parts) > 1 else 0
            return frame_number
        except:
            pass

        # Sometimes beginning of message is sent twice (screen_70screen_707#)
        if len(parts) > 2 and parts[2].isdigit():
            tentative_frame = int(parts[2])
            if self.last_frame - 10 < tentative_frame < self.last_frame + 10:
                return tentative_frame

        # Otherwise trying to make sense of digits
        else:
            digits = ''.join(c for c in ''.join(parts[1:]) if c.isdigit())
            tentative_frame = int(digits) if len(digits) > 1 else 0
            if self.last_frame - 10 < tentative_frame < self.last_frame + 10:
                return tentative_frame

        # Unable to parse - Likely an invalid message
        return None

    def _process_pipe_message(self, message):
        # Parsing
        parts = message.split('#')
        header = parts[0] if len(parts) > 0 else ''
        data = parts[1] if len(parts) > 1 else ''
        parts = header.split('_')
        message_type = parts[0] if len(parts) > 0 else ''
        frame_number = self._parse_frame_number(parts)

        # Invalid message - Ignoring
        if frame_number is None:
            return

        # Processing
        if 'data' == message_type:
            self._process_data_message(frame_number, data)
        elif 'screen' == message_type:
            self._process_screen_message(frame_number, data)
        elif 'tiles' == message_type:
            self._process_tiles_message(frame_number, data)
        elif 'ready' == message_type:
            self._process_ready_message(frame_number)
        elif 'done' == message_type:
            self._process_done_message(frame_number)
        elif 'reset' == message_type:
            self._process_reset_message()
        elif 'exit' == message_type:
            self._process_exit_message()

    def _get_reward(self):
        return self.reward

    def _get_episode_reward(self):
        return self.episode_reward

    def _get_is_finished(self):
        return self.is_finished

    def _get_state(self):
        if 1 == self.draw_tiles:
            return self.tiles.copy()
        else:
            return self.screen.copy()

    def _get_info(self):
        return self.info

    def _reset_info_vars(self):
        self.info = {
            'level': self.level,
            'distance': 0,
            'score': -1,
            'coins': -1,
            'time': -1,
            'player_status': -1
        }

    #Override the step function so that it ignores total reward and episode reward
    def _step(self, action):

        # Changing level
        if self.find_new_level:
            self.change_level()

        obs, step_reward, is_finished, info = NesEnv._step(self, action)
        self.total_reward += step_reward
        self.episode_reward += step_reward
        #reward, self.total_reward = self._calculate_reward(self._get_episode_reward(), self.total_reward)
        # First step() after new episode returns the entire total reward
        # because stats_recorder resets the episode score to 0 after reset() is called
        if self.is_new_episode:
            step_reward = 0.
            #self.total_reward = 0.
            #reward = self.total_reward

        self.is_new_episode = False
        info["level"] = self.level
        info["scores"] = self.get_scores()
        info["total_reward"] = round(self.total_reward, 4)
        info["locked_levels"] = self.locked_levels

        # Indicating new level required
        if is_finished:
            self._unlock_levels()
            self.find_new_level = True

        return obs, step_reward, is_finished, info

class MetaSuperMarioBrosEnv(SuperMarioBrosEnv, MetaNesEnv):

    def __init__(self, average_over=10, passing_grade=600, min_tries_for_avg=5, draw_tiles=0):
        MetaNesEnv.__init__(self,
                            average_over=average_over,
                            passing_grade=passing_grade,
                            min_tries_for_avg=min_tries_for_avg,
                            num_levels=32)
        SuperMarioBrosEnv.__init__(self, draw_tiles=draw_tiles, level=0)
        self.launch_vars['meta'] = '1'

    def _process_reset_message(self):
        self.last_frame = 0

    def _get_standard_reward(self, episode_reward):
        #Returns a standardized reward for an episode (i.e. between 0 and 1,000)
        min_score = 0
        target_score = float(SMB_LEVELS[self.level][MAX_DISTANCE]) - 40
        max_score = min_score + (target_score - min_score) / 0.99  # Target is 99th percentile (Scale 0-1000)
        std_reward = round(1000 * (episode_reward - min_score) / (max_score - min_score), 4)
        std_reward = min(1000, std_reward)  # Cannot be more than 1,000
        std_reward = max(0, std_reward)  # Cannot be less than 0
        return std_reward


import matplotlib.pyplot as plt
import numpy as np
from sklearn import datasets, linear_model
import Myfun
Myfun.pyplot_中文()

# Load the diabetes dataset
diabetes = datasets.load_diabetes()

# Use only one feature
diabetes_X = diabetes.data[:, np.newaxis, 2]  # BMI

# Split the data into training/testing sets
diabetes_X_train = diabetes_X[:-20]
diabetes_X_test = diabetes_X[-20:]

# Split the targets into training/testing sets
diabetes_y_train = diabetes.target[:-20]
diabetes_y_test = diabetes.target[-20:]

# Create linear regression object
regr = linear_model.LinearRegression()

# Train the model using the training sets
regr.fit(diabetes_X_train, diabetes_y_train)

# The coefficients
print('Coefficients: \n', regr.coef_)
# The mean squared error
print("Mean squared error: %.2f"
      % np.mean((regr.predict(diabetes_X_test) - diabetes_y_test) ** 2))
# Explained variance score: 1 is perfect prediction
print('Variance score: %.2f' % regr.score(diabetes_X_test, diabetes_y_test))

# Plot outputs

plt.scatter(diabetes_X_test, diabetes_y_test,  color='black')
plt.plot(diabetes_X_test, regr.predict(diabetes_X_test), color='blue',
         linewidth=3)
plt.xlabel("BMI")
plt.ylabel("Target 一年後的患病指數")
# plt.xticks(())
# plt.yticks(())

plt.show()

Regression回歸分析/糖尿病/回歸分析找出BMI與糖尿病關係.py

回歸分析找出BMI與糖尿病關係.py


"""This module implements the following classes:
    - AnimalsDataListView
"""

from PyQt5 import QtCore, QtGui, QtWidgets

from pigcel.gui.models.animals_data_model import AnimalsDataModel
from pigcel.gui.views.double_clickable_listview import DoubleClickableListView

class AnimalsDataListView(DoubleClickableListView):
    """This class implments the listview used to drop in animal names.
    """

    def keyPressEvent(self, event):
        """Event handler for keyboard interaction.

        Args:
            event (PyQt5.QtGui.QKeyEvent): the keyboard event
        """

        if event.key() == QtCore.Qt.Key_Delete:

            animals_data_model = self.model()
            if animals_data_model is None:
                return

            for sel_index in reversed(self.selectedIndexes()):
                animals_data_model.remove_workbook(sel_index.data(AnimalsDataModel.Workbook))
                if animals_data_model.rowCount() > 0:
                    index = animals_data_model.index(animals_data_model.rowCount()-1)
                    self.setCurrentIndex(index)

        else:
            super(AnimalsDataListView, self).keyPressEvent(event)


"""
Developed By: JumpShot Team
Written by: riscyseven
"""

from PyQt6.QtWidgets import QDialog, QMessageBox
from PyQt6 import uic
from editor.proptab.fileseldialog import FileSelWidget

from gm_resources import resourcePath
from attr import CompAttr

class GlobalDirectory(QDialog):
    def __init__(self, dirName, parent=None):
        super().__init__(parent)
        path = resourcePath("src/editor/globaldirectory.ui")
        uic.loadUi(path, self) # Load the .ui file
        self.fileDialog = FileSelWidget(CompAttr.DRSAVE, self._fileSelected, dirName, self)
        self.verticalLayout.addWidget(self.fileDialog)
        self.setWindowTitle("Set Global File Output Location")
        self.pushButton.clicked.connect(self._setClicked)
        self.dir = dirName

    def _fileSelected(self, dirName):
        self.dir = dirName

    def _setClicked(self):
        if (QMessageBox.warning(self, "Are you sure?", "This action will:\n"
            "   1. Change the root directory of each display component\n"
            "   2. Change all output file names to its component names\n"
            "Do you want to proceed?", QMessageBox.StandardButton.Ok |
                 QMessageBox.StandardButton.Cancel) == QMessageBox.StandardButton.Cancel):
                 self.dir = ""
        self.close()

    def getDirectory(self) -> str:
        return self.dir


import requests
from Bio import SeqIO
import subprocess
import urllib.request
import sys
import os
import os.path

def requests_from_uniprot(uniprot_id, directory):
    '''
    Function takes Uniprot ID \n
    Return file.fasta of this target protein

    :param uniprot_id: id of target protein
    :type uniprot_id: str
    :param directory: dir  where the fasta file  will be saved
    :type directory: str
    :return: fasta file with target protein
    :rtype: fasta.file
    '''
    base = 'http://www.uniprot.org'
    kb_endpoint = '/uniprot/'
    payload = {'query': f'{uniprot_id}',
               'format': 'fasta'}
    result = requests.get(base + kb_endpoint, params=payload)
    if result.ok:
        with open(os.path.join(directory, f"{uniprot_id}.fasta"), "w") as f:
            f.write(result.text)
    else:
        return print('Something went wrong: ', result.status_code)

def clean_db_pdb(name, directory):
    '''
    This function cleaning wwPDB database from duplicates

    :param name: name fasta file
    :type name: str
    :param directory: path to folder where will be saved file
    :type directory: str
    :return: clear database
    :rtype: file
    '''
    fasta_sequences = SeqIO.parse(open(name), 'fasta')
    first_elem = str()
    with open(os.path.join(directory, "clean_pdb_db.fasta"), "w") as f:
        for fasta in fasta_sequences:
            seq = fasta.seq
            if seq != first_elem:
                first_elem = seq
                SeqIO.write(fasta, f, 'fasta')

def downloading_databases(directory, only_swiss_model=False, only_wwpdb_model=False, clean_pdb_db=True):
    '''
    Function take path to directory and flag for download db. In this function user can choose which db will be used
    for search homologous. The flags are used to select the db.

    :param directory: dir where the db file will be saved
    :type directory: str
    :param only_swiss_model: if choose this flag only swiss_model.db will be download and will be used for search of
    homologous
    :type only_swiss_model: bool
    :param only_wwpdb_model: if choose this flag wwpdb.db will be download and will be used for search of
    homologous
    :type only_wwpdb_model: bool
    :param clean_pdb_db: if choose this flag wwpdb.db will be cleaned from duplicates with using function clean_db_pdb
    :type clean_pdb_db: bool
    :return: file with merge db or not merge db
    :rtype: file
    '''
    out_uniprot = os.path.join(directory, "uniprot_sprot.fasta.gz")
    out_pdb_db = os.path.join(directory, "pdb_seqres.fasta")
    new_name_out_uniprot = os.path.join(directory, "uniprot_sprot.fasta")
    out_protein_db = os.path.join(directory, "db_for_protein.fasta")
    if only_swiss_model:
        urllib.request.urlretrieve(
            "ftp://ftp.uniprot.org/pub/databases/uniprot/current_release/knowledgebase/complete/uniprot_sprot.fasta.gz",
            out_uniprot)
        gunzip_uniprot = subprocess.Popen([f"gunzip {out_uniprot}"], shell=True, encoding='utf-8')
        gunzip_uniprot = gunzip_uniprot.wait()
        subprocess.Popen([f"mv {new_name_out_uniprot} {out_protein_db}"], shell=True, encoding='utf-8')
    elif only_wwpdb_model:
        urllib.request.urlretrieve("https://ftp.wwpdb.org/pub/pdb/derived_data/pdb_seqres.txt", out_pdb_db)
        if clean_pdb_db:
            clean_db_pdb(out_pdb_db, directory)
            new_name_pdb_db = os.path.join(directory, "clean_pdb_db.fasta")
            subprocess.Popen([f"mv {new_name_pdb_db} {out_protein_db}"], shell=True, encoding='utf-8')
        else:
            subprocess.Popen([f"mv {out_protein_db} {out_protein_db}"], shell=True, encoding='utf-8')
    else:
        urllib.request.urlretrieve(
            "ftp://ftp.uniprot.org/pub/databases/uniprot/current_release/knowledgebase/complete/uniprot_sprot.fasta.gz",
            out_uniprot)
        gunzip_uniprot = subprocess.Popen([f"gunzip {out_uniprot}"], shell=True, encoding='utf-8')
        gunzip_uniprot = gunzip_uniprot.wait()
        urllib.request.urlretrieve("https://ftp.wwpdb.org/pub/pdb/derived_data/pdb_seqres.txt", out_pdb_db)
        if clean_pdb_db:
            clean_db_pdb(out_pdb_db, directory)
            new_name_db = os.path.join(directory, "clean_pdb_db.fasta")
            out_protein_db = os.path.join(directory, "db_for_protein.fasta")
            subprocess.Popen([f"cat {new_name_db} {new_name_out_uniprot} > {out_protein_db}"],
                             shell=True, encoding='utf-8')
        else:
            subprocess.Popen([f"cat {out_pdb_db} {new_name_out_uniprot} > {out_protein_db}"],
                             shell=True, encoding='utf-8')
    print("Databases have been downloaded !")

def create_database_for_mafft(directory):
    '''
    Created special database for mafft tools

    :param directory: path to folder
    :type directory: str
    :return: folder with database for mafft
    :rtype: dir
    '''
    name_of_dir_with_db = os.path.join(directory, "db_for_mafft")
    name_protein_db = os.path.join(directory, "db_for_protein.fasta")
    if not os.path.exists(name_of_dir_with_db):
        os.makedirs(name_of_dir_with_db)
    mafft_database = subprocess.Popen(
        [
            f"makeblastdb -in {name_protein_db} -out {name_of_dir_with_db}/db_for_mafft -dbtype prot -parse_seqids -hash_index"],
        shell=True, encoding='utf-8')
    (output, err) = mafft_database.communicate()
    mafft_database = mafft_database.wait()
    if err:
        print(err)
    print("Database for mafft created")

def start_mafft(directory, id_protein, n=30, th=40, flag_path=False, path=None):
    '''
    Start mafft with L-INS-i algorithm
    (more information here:https://mafft.cbrc.jp/alignment/software/algorithms/algorithms.html)
    User can choose number of homologoues and level of threshold line.

    :param directory: dir where the db file will be saved
    :type directory: str
    :param id_protein: fasta file with target protein
    :type id_protein: file
    :param n: number of homologoues
    :type n: int
    :param th: of threshold line
    :type th: int
    :param flag_path: if user have db which was created earlier takes this flag
    :type flag_path: bool
    :param path: path to db which was created earlier
    :type path: bool
    :return: file_output.fasta with list of aligned homologues
    :rtype: file
    '''
    name_of_dir_with_db = os.path.join(directory, "db_for_mafft/db_for_mafft")
    path_to_fasta = os.path.join(directory, f"{id_protein}.fasta")
    path_to_out_fasta = os.path.join(directory, f"{id_protein}_output.fasta")
    if flag_path == False:
        path = name_of_dir_with_db
    mafft_homologs = subprocess.Popen([f"mafft-homologs -l -d {path} \
                                        -o '--thread 8 --threadtb 5 --threadit 0 --reorder\
                                        --anysymbol --maxiterate 1000 --retree 1 --localpair' \
                                        -a {n} -e 1.0e-{th} -f {path_to_fasta} > {path_to_out_fasta}"],
                                      shell=True,
                                      encoding='utf-8')
    (output, err) = mafft_homologs.communicate()
    if err:
        print(err)
    mafft_homologs = mafft_homologs.wait()
    print(f"Mafft searched {n} homologues")

def fasta_parser(directory, name_fasta, name_with_chain=False):
    '''
    Easy label of protein fasta parser

    :param directory: dir where the db file will be saved
    :type directory: str
    :param name_fasta: file_output.fasta with list of aligned homologues
    :type name_fasta: file
    :param name_with_chain: choose this flag if user want to save list with name which not only name but name with chain
    :type name_with_chain: bool
    :return: list with id name of protein homologues
    :rtype: list
    '''
    fasta_name_and_path = os.path.join(directory, f"{name_fasta}_output.fasta")
    fasta_sequences = SeqIO.parse(open(fasta_name_and_path), 'fasta')
    for fasta in fasta_sequences:
        break
    name_list = []
    for fasta in fasta_sequences:
        if name_with_chain:
            fasta_name = fasta.id.split("|")
            if fasta_name[1][0].isdigit():
                fasta_name = fasta_name[1] + "_" + fasta_name[2]
            else:
                fasta_name = fasta_name[1]
            name_list.append(fasta_name)
        else:
            name_list.append(fasta.id.split("|")[1])
    return name_list

def pdb_download(directory, name_fasta):
    '''
    Download all homologues in pdb format \n
    Returns files.pdb

    :param directory: dir where the db file will be saved
    :type directory: str
    :param name_fasta: file_output.fasta with list of aligned homologues
    :type name_fasta: file
    :return: dir with downloaded pdb structure
    :rtype: file
    '''
    path_to_orig_templates = os.path.join(directory, "orig_templates")
    if not os.path.exists(path_to_orig_templates):
        os.makedirs(path_to_orig_templates)
    list_pdb_structure = fasta_parser(directory, name_fasta)
    list_with_name = fasta_parser(directory, name_fasta, name_with_chain=True)
    for name in range(0,len(list_pdb_structure)):
        if list_pdb_structure[name][0].isdigit():
            pdbfn = list_pdb_structure[name] + ".pdb"
            name_with_chain = list_with_name[name] + ".pdb"
            url = "https://files.rcsb.org/download/" + pdbfn
            outfnm = os.path.join(path_to_orig_templates, name_with_chain)
            try:
                urllib.request.urlretrieve(url, outfnm)
                print(f"Structure {name_with_chain} successfully downloaded")
            except Exception as err:
                print(str(err), file=sys.stderr)
        else:
            uniname = list_pdb_structure[name] + ".pdb"
            url_uni = "https://swissmodel.expasy.org/repository/uniprot/" + uniname
            outfnm_uni = os.path.join(path_to_orig_templates, uniname)
            try:
                urllib.request.urlretrieve(url_uni, outfnm_uni)
                print(f"Structure {uniname} successfully downloaded")
            except Exception as err:
                print(str(err), file=sys.stderr)


__copyright__ = "Copyright (c) 2020 Jina AI Limited. All rights reserved."
__license__ = "Apache-2.0"

import numpy as np

from .. import BaseAudioEncoder
from ...decorators import batching, as_ndarray

class MFCCTimbreEncoder(BaseAudioEncoder):
    """
    :class:`MFCCTimbreEncoder` is based on Mel-Frequency Cepstral Coefficients (MFCCs) which represent timbral features.
    :class:`MFCCTimbreEncoder` encodes an audio signal from a `Batch x Signal Length` ndarray into a
    `Batch x Concatenated Features` ndarray.
    """

    def __init__(self, input_sample_rate: int = 22050, n_mfcc: int = 20, n_fft_length: int = 2048,
                 hop_length: int = 512, *args, **kwargs):
        """
        :class:`MFCCTimbreEncoder` extracts from an audio signal a `n_mfcc`-dimensional feature vector for each MFCC
        frame.

        :param input_sample_rate: input sampling rate in Hz (22050 by default)
        :param n_mfcc: the number of coefficients (20 by default)
        :param n_fft: length of the FFT window (2048 by default)
        :param hop_length: the number of samples between successive MFCC frames (512 by default)
        """
        super().__init__(*args, **kwargs)
        self.input_sample_rate = input_sample_rate
        self.n_mfcc = n_mfcc
        self.n_fft_length = n_fft_length
        self.hop_length = hop_length

    @batching
    @as_ndarray
    def encode(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        """
        Segments the audio signal of each Chunk into short MFCC frames, extracts MFCCs for each frame and concatenates
        Chunk frame MFCCs into a single Chunk embedding.

        :param data: a `Batch x Signal Length` ndarray, where `Signal Length` is a number of samples
        :return: a `Batch x Concatenated Features` ndarray, where `Concatinated Features` is a `n_mfcc`-dimensional
        feature vector times the number of the MFCC frames
        """
        from librosa.feature import mfcc
        embeds = []
        for chunk_data in data:
            mfccs = mfcc(y=chunk_data, sr=self.input_sample_rate, n_mfcc=self.n_mfcc, n_fft=self.n_fft_length,
                         hop_length=self.hop_length)
            embeds.append(mfccs.flatten())
        return embeds

class ChromaPitchEncoder(BaseAudioEncoder):
    """
    :class:`ChromaPitchEncoder` is based on chroma spectrograms (chromagrams) which represent melodic/harmonic features.
    :class:`ChromaPitchEncoder` encodes an audio signal from a `Batch x Signal Length` ndarray into a
    `Batch x Concatenated Features` ndarray.
    """

    def __init__(self, input_sample_rate: int = 22050, hop_length: int = 512, *args, **kwargs):
        """
        :class:`ChromaPitchEncoder` extracts from an audio signal a 12-dimensional feature vector, that refer to 12
        octaves, for each chroma frame.

        :param input_sample_rate: input sampling rate in Hz (22050 by default)
        :param hop_length: the number of samples between successive chroma frames (512 by default)
        """
        super().__init__(*args, **kwargs)
        self.input_sample_rate = input_sample_rate
        self.hop_length = hop_length

    @batching
    @as_ndarray
    def encode(self, data: np.ndarray, *args, **kwargs) -> np.ndarray:
        """
        Segments the audio signal of each Chunk into short chroma frames, extracts chromagrams for each frame and
        concatenates Chunk frame chromagrams into a single Chunk embedding.

        :param data: a `Batch x Signal Length` ndarray, where `Signal Length` is a number of samples
        :return: a `Batch x Concatenated Features` ndarray, where `Concatinated Features` is a 12-dimensional feature
        vector times the number of the chroma frames
        """
        from librosa.feature import chroma_cqt
        embeds = []
        for chunk_data in data:
            chromagrams = chroma_cqt(y=chunk_data, sr=self.input_sample_rate, n_chroma=12, hop_length=self.hop_length)
            embeds.append(chromagrams.flatten())
        return embeds