     
     // MARK: - Lifecycle
   
-    public init(staticCellHeight: CGFloat, nextLayoutStaticCellHeight: CGFloat, layoutState: LayoutState) {
+    public init(staticCellHeight: CGFloat, nextLayoutStaticCellHeight: CGFloat, layoutState: LayoutState, cellHeightPadding: CGFloat = 6, cellWidthPadding: CGFloat = 6, gridLayoutCountOfColumns: Int = 3) {
         self.staticCellHeight = staticCellHeight
-        self.numberOfColumns = layoutState == .list ? ListLayoutCountOfColumns : GridLayoutCountOfColumns
+        self.numberOfColumns = layoutState == .list ? listLayoutCountOfColumns : gridLayoutCountOfColumns
         self.layoutState = layoutState
         self.nextLayoutStaticCellHeight = nextLayoutStaticCellHeight
+        self.cellHeightPadding = cellHeightPadding
+        self.cellWidthPadding = cellWidthPadding
+        self.gridLayoutCountOfColumns = gridLayoutCountOfColumns
         
         super.init()
     }

open class DisplaySwitchLayout: UICollectionViewLayout {
    
    fileprivate let numberOfColumns: Int
    fileprivate let cellPadding: CGFloat = 6.0
    fileprivate let staticCellHeight: CGFloat
    fileprivate let nextLayoutStaticCellHeight: CGFloat
    fileprivate var previousContentOffset: NSValue?
    fileprivate var layoutState: LayoutState
  
    fileprivate var baseLayoutAttributes: [DisplaySwitchLayoutAttributes]!
    

        var yOffset = [CGFloat](repeating: contentHeight, count: numberOfColumns)
        for item in 0 ..< collectionView!.numberOfItems(inSection: 0) {
            let indexPath = IndexPath(item: item, section: 0)
            let height = cellHeightPadding + staticCellHeight
            let frame = CGRect(x: xOffsets[column], y: yOffset[column], width: columnWidth, height: height)
            let insetFrame = frame.insetBy(dx: cellWidthPadding, dy: cellHeightPadding)
            let attributes = DisplaySwitchLayoutAttributes(forCellWith: indexPath)
            attributes.frame = insetFrame
            baseLayoutAttributes.append(attributes)
