from ..pakbase import Package


class ModflowPbc(Package):
    """
    Periodic boundary condition class

    """

    def __init__(self, model, layer_row_column_data=None,
                 layer_row_column_shead_ehead=None,
                 cosines=None, extension='pbc', unitnumber=None,
                 zerobase=True):
        # set default unit number of one is not specified
        if unitnumber is None:
            unitnumber = ModflowPbc.defaultunit()

        # Call ancestor's init to set self.parent, extension, name and
        # unit number
        Package.__init__(self, model, extension, ModflowPbc.ftype(),
                         unitnumber)
        self.heading = '# {} package for '.format(self.name[0]) + \
                       ' {}, '.format(model.version_types[model.version]) + \
                       'generated by Flopy.'
        self.mxactp = 0
        if layer_row_column_data is None:
            if layer_row_column_shead_ehead is not None:
                msg = '\nWARNING: ModflowPbc - Do not use ' + \
                      'layer_row_column_shead_ehead!\n' + \
                      22 * ' ' + 'Use layer_row_column_data instead.'
                print(msg)
                layer_row_column_data = layer_row_column_shead_ehead
            else:
                e = 'Failed to specify layer_row_column_shead_ehead ' + \
                    'or layer_row_column_data.'
                raise Exception(e)

        self.mxactp, self.layer_row_column_data = self.assign_layer_row_column_data(
            layer_row_column_data, 5, zerobase=zerobase)
        # misuse of this function - zerobase needs to be False
        self.mxcos, self.cosines = self.assign_layer_row_column_data(cosines,
                                                                     3,
                                                                     zerobase=False)
        # self.mxcos = 0
        # if (cosines != None):
        #     error_message = 'cosines must have 3 columns'
        #     if (not isinstance(cosines, list)):
        #         cosines = [cosines]
        #     for a in cosines:
        #         a = np.atleast_2d(a)
        #         nr, nc = a.shape
        #         assert nc == 3, error_message
        #         if (nr > self.mxcos):
        #             self.mxcos = nr
        #     self.cosines = cosines
        self.np = 0
        self.parent.add_package(self)

    def ncells(self):
        # Returns the  maximum number of cells that have recharge
        # (developed for MT3DMS SSM package)
        return self.mxactp

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

        Returns
        -------
        None

        """
        f_pbc = open(self.fn_path, 'w')
        f_pbc.write('%s\n' % self.heading)
        f_pbc.write('%10i%10i\n' % (self.mxactp, self.mxcos))
        for n in range(self.parent.get_package('DIS').nper):
            if (n < len(self.layer_row_column_data)):
                a = self.layer_row_column_data[n]
                itmp = a.shape[0]
            else:
                itmp = -1
            if (n < len(self.cosines)):
                c = self.cosines[n]
                ctmp = c.shape[0]
            else:
                ctmp = -1
            f_pbc.write('{:10d}{:10d}{:10d}\n'.format(itmp, ctmp, self.np))
            if n < len(self.layer_row_column_data):
                for b in a:
                    line = '{:10d}{:10d}{:10d}{:10d}{:10d}\n'.format(b[0],
                                                                     b[1],
                                                                     b[2],
                                                                     b[3],
                                                                     b[4])
                    f_pbc.write(line)
            if n < len(self.cosines):
                for d in c:
                    f_pbc.write('{:10g}{:10g}{:10g}\n'.format(d[0],
                                                              d[1],
                                                              d[2]))
        f_pbc.close()

    @staticmethod
    def ftype():
        return 'PBC'

    @staticmethod
    def defaultunit():
        return 30
