from ..pakbase import Package


class Mt3dTob(Package):
    """
    Transport Observation package class
    """

    def __init__(
        self,
        model,
        outnam="tob_output",
        CScale=1.0,
        FluxGroups=[],
        FScale=1.0,
        iOutFlux=0,
        extension="tob",
        unitnumber=None,
        filenames=None,
    ):
        if unitnumber is None:
            unitnumber = Mt3dTob._defaultunit()
        elif unitnumber == 0:
            unitnumber = Mt3dTob._reservedunit()

        # call base package constructor
        super().__init__(
            model,
            extension=extension,
            name=self._ftype(),
            unit_number=unitnumber,
            filenames=self._prepare_filenames(filenames),
        )

        self.heading = "# TOB for MT3DMS, generated by Flopy."
        self.outnam = outnam
        self.CScale = CScale
        self.FluxGroups = FluxGroups
        self.FScale = FScale
        self.iOutFlux = iOutFlux
        self.parent.add_package(self)
        return

    def __repr__(self):
        return "Transport Observation package class"

    def write_file(self):
        """
        Write the package file

        Returns
        -------
        None

        """
        # Open file for writing
        f_tob = open(self.fn_path, "w")
        f_tob.write("%s\n" % (self.heading))
        MaxConcObs = 0
        MaxFluxObs = 0
        MaxFluxCells = 0
        inConcObs = 0
        inFluxObs = 88
        inSaveObs = 89
        if inFluxObs:
            for FluxGroup in self.FluxGroups:
                MaxFluxCells = MaxFluxCells + len(FluxGroup[1])
                MaxFluxObs = MaxFluxObs + 1
        f_tob.write("%10d%10d%10d\n" % (MaxConcObs, MaxFluxObs, MaxFluxCells))
        f_tob.write("%s%10d%10d%10d\n" % (self.outnam, inConcObs, inFluxObs, inSaveObs))

        if inFluxObs:
            nFluxGroup = len(self.FluxGroups)
            f_tob.write("%10d%10f%10d\n" % (nFluxGroup, self.FScale, self.iOutFlux))
            for FluxGroup in self.FluxGroups:
                nFluxTimeObs, FluxTimeObs = self.assign_layer_row_column_data(
                    FluxGroup[0], 5, zerobase=False
                )  # misuse of function - zerobase set to False
                nCells, Cells = self.assign_layer_row_column_data(
                    FluxGroup[1], 4, zerobase=False
                )  # misuse of function - zerobase set to False
                nCells = 4
                iSSType = FluxGroup[2]
                f_tob.write("%10d%10d%10d\n" % (nFluxTimeObs, nCells, iSSType))
                for fto in FluxTimeObs:
                    fto = fto[0]  # Still to fix this!
                    f_tob.write(
                        "%12s%10s%10s%10s%10s\n"
                        % (fto[0], fto[1], fto[2], fto[3], fto[4])
                    )
                for c in Cells:
                    c = c[0]  # Still to fix this!
                    f_tob.write("%10d%10d%10d%10f\n" % (c[0], c[1], c[2], c[3]))

        f_tob.close()
        return

    @staticmethod
    def _ftype():
        return "TOB"

    @staticmethod
    def _defaultunit():
        return 37

    @staticmethod
    def _reservedunit():
        return 12
